[KMCompiler][MetaX][Hygon][Mthreads][T-Head][Ascend] fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert backends, plus test and benchmark fixes

[cheersluvs]

Backend overrides of fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert for
MetaX C550, Hygon BW1000, Moore Threads S5000, T-Head PPU-ZW810E and Ascend
910B, plus seven fixes to this operator's test and benchmark that a card other
than an 80 GiB H800 exposes. Twelve files, all belonging to this operator — no
shared module added or changed.

The five overrides are not the same change. Three are token tiling (C550,
BW1000, S5000) with per-card launch constants and dispatch rules; one is an
enablement, not a tuning (ZW810E, where the operator could not be built at all);
and one is both (910B, which needed eight compiler workarounds before it ran and
then eight tuning steps).

One of the seven fixes is a prerequisite for all of it: as the file stands, a
vendor override of this operator can never be reached from its own test or
benchmark. Details in the last section.

Performance against achievable bandwidth

The ceiling is taken warm, which matters more than it sounds. On C550 a
512 MiB device-to-device copy reads 1209.7 GB/s with three warmup iterations and
1396.6 / 1397.7 / 1397.9 once the clocks have ramped; measuring it cold at the
top of a script inflates every percentage computed against it — enough, in one
draft of this table, to put the operator above 100% of its own ceiling. Only
MetaX ramps: the same probe on BW1000 gives 1343.9 cold against 1350.1 / 1349.9
/ 1349.9 warm, and on S5000 the ceiling taken after a full load is within 0.3% of
the one taken before it. Every card here uses the warm figure.

	ceiling	generic	override
MetaX C550, 131072 x 64	1397 GB/s	1162.8 (83.2%)	1337.6 (95.7%)
MetaX C550, 32768 x 128	1397 GB/s	1224.2 (87.6%)	1355.5 (97.0%)
Hygon BW1000, 32768 x 64	1350 GB/s	609.9 (45.2%)	1197.8 (88.7%)
Hygon BW1000, 32768 x 128	1350 GB/s	611.6 (45.3%)	1209.6 (89.6%)
Hygon BW1000, 131072 x 128	1350 GB/s	612.3 (45.4%)	1213.8 (89.9%)
MTT S5000, 32768 x 64	1332 GB/s	1162.5 (87.4%)	1220.3 (91.6%)
MTT S5000, 32768 x 128	1332 GB/s	1245.0 (93.6%)	no config wins, falls through
T-Head ZW810E, 32768 x 64	2097 GB/s	cannot be built	2013.3 (96.0%)
T-Head ZW810E, 131072 x 128	2097 GB/s	cannot be built	2014.1 (96.0%)
Ascend 910B, 32768 x 64	1229 GB/s	cannot be built	179.2 (14.6%)
Ascend 910B, 32768 x 128	1229 GB/s	cannot be built	265.1 (21.6%)

Ascend is the outlier and the percentage is the honest number, not a typo. That
card has no tl.float8e4nv, no pointer bitcast, no scalar float-to-int and a
65535-program launch cap, so there is no generic build to compare against; what
the override buys is measured below as a 6.7x / 9.9x improvement over the
first version that ran at all. The 1229 GB/s is the low end of a 1229-1433 GB/s
copy ceiling. Do not read the 14.6% as headroom that a day of tuning would take
— the closing ledger below lists what was tried.

Performance against a baseline

Through this repo's benchmark, --mode kernel --level core, 22 shapes, except
where noted. Which baseline is available differs per card and each one's
provenance is stated after the table.

	64 heads	128 heads
C550 vs mcoplib, 8192-131072	0.870-0.888 → 1.004-1.015	0.945-0.951 → 1.047-1.052
C550 vs mcoplib, 1-64 tokens	1.23-1.61x, generic path both ways	1.27-1.52x, generic path both ways
BW1000 vs vLLM XPU Triton, 8192-131072	2.02-2.04x	2.99-3.22x
BW1000 vs vLLM XPU Triton, 1024-2048	1.94-1.98x	1.96-2.30x
BW1000 vs vLLM XPU Triton, 1-17 tokens	2.99-3.63x	2.47-3.62x
S5000 vs vLLM XPU Triton, 8192-131072	1.144-1.208 → 1.205-1.264	1.087-1.150, generic path
S5000 vs vLLM XPU Triton, 1-64 tokens	1.38-1.73x, generic path both ways	1.29-1.88x, generic path both ways
ZW810E vs T-Head's own vLLM, 32768-131072	1.24-1.25x	1.21-1.22x
ZW810E vs T-Head's own vLLM, 1024-8192	1.25-1.32x	1.23-1.26x
ZW810E vs T-Head's own vLLM, 1-64 tokens	1.34-1.41x	1.32-1.41x
910B vs an eager torch_npu composition, 32768-131072	11.3-11.8x	10.8-11.0x (to 65536)
910B vs an eager torch_npu composition, 1-64 tokens	7.6-8.8x	8.5-9.9x

The C550 rows need a disclosure, because the baseline cannot run unaided.
MetaX's kernel returns mcErrorInvalidValue from every launch of this op on C550
— every shape, both dtypes, with a valid configuration — so with a stock
environment this benchmark reports 1 skipped, which is what the wrapper
described further down is for. The numbers above were taken with an LD_PRELOAD
shim that forwards wcudaLaunchKernelExC to the classic mcLaunchKernel. That
forwarding is semantically exact here: numAttrs == 0 on these launches, so the
extended call carries nothing the classic one lacks — grid, block, shared memory
and stream. And the forwarded baseline was checked rather than assumed: running
test_matches_vllm_reference under the same shim on the same card passes 16/16 at
the accuracy suite's own tolerances, which rules out the specific failure mode
worth worrying about (MetaX's launcher computes kWarpsPerBlock = kBlockSize / 32
while C550 warps are 64 wide, so a half-covered grid would have looked
impressively fast). Reported as
MetaX-MACA/mcoplib#59. Both
halves of the C550 row come from the same card, branch, shim and harness — the
only difference is whether the override is registered.

The BW1000 and S5000 baselines need their provenance stated: they are the same
code, patched, and it is local-only. Neither card has a vendor kernel for this
operator. On BW1000 the installed vLLM (0.18.1+das.fa71803.dtk2604) does not
ship the operator in any form — searching its tree for deepseek_v4 or
xpu_qnorm_rope returns nothing; on S5000 the installed vLLM is an .empty
build registering 0 deepseek ops, and vllm-musa 0.1.1 is a platform plugin,
not an operator library. Run as-is this benchmark reports 1 skipped on both.
The baseline above is vLLM's own portable XPU Triton implementation taken from
v0.27.1, checked function by function against upstream rather than trusted:
quantize_and_insert_k_kernel is byte-identical, and _xpu_qnorm_rope_kernel
differs by exactly six lines, which are two one-line fixes — tl.program_id cast
to int64, and a mask=offs < NOPE_DIM on the KV bulk store. The host wrappers
additionally drop vLLM-internal imports and the FNUZ branch (it needs
vllm.platforms) and gain a num_warps argument. None of this is in the
diff — it is wired in by hand for the measurement and reverted afterwards.

Both fixes are needed to quote anything at these sizes: without the cast the
address computation wraps past 2^31 q elements, which 65536 x 64 already reaches;
without the mask the bulk store writes the unrotated RoPE region and a later store
overwrites it with no ordering, corrupting ~0.7% of tokens nondeterministically.
Both also make the baseline faster — the mask cuts total stores per token from
576 to 512 elements — so patching it makes this comparison harder on us, not
easier. Reported upstream as
vllm-project/vllm#52415 and
#52416.

The baseline runs at its own measured optimum on each card, not at a value
carried over. On BW1000 that is num_warps=1 (590.6 versus 572.9 / 461.8 / 241.6
GB/s at 2 / 4 / 8); on S5000 it is also num_warps=1, at 982.5 GB/s against
445.6 at num_warps=4. That second figure is worth stating because an earlier
S5000 table quoted the num_warps=4 number: the baseline was crippled 2.2x —
worse than this operator was under the same compiler defect — so the old table
understated this operator, and the apparent 64-head "parity" in it was the
coincidence of both implementations landing near 446 GB/s. A benchmark harness
will not tell you a baseline is handicapped; back-compute GB/s from its own
reported latency and compare against a known bandwidth before believing any ratio.

What "as shipped" would give, and why this PR does not quote it

Upstream passes no num_warps — v0.27.1 and main both, at
vllm/models/deepseek_v4/xpu/xpu_qnorm_rope_kv_fp8_insert.py:135, with no
@triton.autotune — so running the baseline verbatim takes Triton's default of
4. Same harness, same 22 shapes, baseline rebuilt from upstream with both file
sha256 verified and only the two portability patches above applied:

8192-131072	as shipped (num_warps absent → 4)	at the baseline's own optimum
BW1000, 64 heads	2.57-2.62x	2.02-2.04x
BW1000, 128 heads	2.89-3.03x	2.99-3.22x
S5000, 64 heads	2.68-2.81x	1.205-1.264x
S5000, 128 heads	2.81-2.92x	1.087-1.150x

On BW1000 the launch configuration only matters at 64 heads. There the
default takes the baseline from 45% to 34.5% of achievable and the ratio from
2.02 to 2.62. At 128 heads the two columns overlap and the baseline sits at
~30% of achievable either way: it is already degraded there for some reason that
is not the launch configuration, and this PR does not explain that — it is the
same unexplained 128-head effect noted in the decomposition above, now with one
more property rather than a mechanism.

On S5000 the default costs the baseline 982.5 → 445.6 GB/s and inflates the
ratio 2.3x across both head counts.

Two checks on the as-shipped run itself. It reproduces the num_warps=4 point
of the BW1000 sweep quoted above independently — 461.4 GB/s measured at
32768 x 64 against 461.8 in that sweep — and on both cards the as-shipped
baseline lands at 33-34% of a ceiling it otherwise reaches 45% of, which is the
signature of a handicapped launch rather than a slow kernel.

This PR quotes the at-optimum column. A gap that one launch argument closes
is not this operator's contribution. Rows below ~192 tokens are omitted from the
as-shipped table: that region spreads 12-34% across repetitions on these parts,
and the as-shipped run was single-pass.

The BW1000 45% → 90% figures were measured separately on this branch by comparing
flaggems_vllm.ops.<op> against flaggems_vllm.<op> in one process — the generic
implementation and whatever the vendor registrar bound — so there is no file
editing between the two halves and no chance of drift.

On BW1000, that there is no vendor kernel was established by symbol scan rather
than by absence of an import: vllm._custom_ops registers 49 ops and none is this
one, and 2505 shared objects export no matching symbol. The vendor's lightop
does cover this stage — fused_rms_norm_rope_contiguous and
fuse_rmsnorm_rope_quant_qkv both take slot_mapping and kv_cache — but only
for a scalar-scale FP8 convention: kv_cache_scale is a single float, which
cannot carry seven per-block exponents per token, so that interface cannot express
this operator's UE8M0 blockwise layout whatever the kernel does internally.

The 910B row is not from this benchmark. The harness gates on
VLLM_REF_AVAILABLE and self-skips there, because no baseline of any kind exists
on that card — not a vendor kernel, and not even vLLM's portable Triton fallback,
which does not compile (details in the Ascend section). The eager torch_npu
composition it is measured against is described there too. Those figures come
from a standalone script, and saying so is the point: they are not the same
measurement as the rows above.

What the token tiling does

This part applies to C550, BW1000 and S5000. The generic kernel runs one program
per (token, head) slot at num_warps=1 — 64 threads for 512 elements on a
64-lane warp, which is 8.5M programs of 64 threads at 131072 x 64. Two costs
follow: the KV slot does seven quantisation blocks against a Q slot's
normalise-and-rotate, so one slot per program makes it a straggler; and 64-thread
blocks under-fill the SM. The straggler is 1-in-65 at 64 heads but 1-in-129 at
128, matching the measured shortfall being worse at the lower head count.

Each program here handles TPP tokens of one slot, so every program is
uniformly all-Q or all-KV with no divergence. On C550 and BW1000 that is
TPP=8/num_warps=4: 16 elements per lane instead of 8, in a 256-thread block
instead of 64.

Two axes matter and neither is visible on its own, which is worth recording
because tuning them separately gives the wrong answer twice over. A full
TPP x num_warps sweep on BW1000 puts every optimum at TPP / num_warps = 2 —
two tokens per warp, so 16 elements per lane. Bandwidth by elements per lane is
237 / 408 / 604 / 906 / 1183 / 1077 / 1050 GB/s for 1 / 2 / 4 / 8 / 16 / 32 /
64. But elements per lane does not explain everything: TPP=1/warps=1 and
TPP=2/warps=2 are both 8 elements per lane and differ by 50% (604 vs 906),
because the second has a wider program. Meanwhile raising num_warps alone at
TPP=1 measurably hurts on C550 (1134 → 691 → 398 → 199 GB/s) and shows 8 and 4
elements per lane tied at ~605 on BW1000 — which invites the conclusion that
access width does not matter. At TPP=1 the block is only 512 elements and there
is nothing to widen into.

Do not read a (TPP, num_warps) pair across cards. S5000 has a 32-lane
warp where the other two have 64, so its stock num_warps=1 launch already
delivers 16 elements per lane — the value C550 and BW1000 need TPP=8/warps=4 to
reach. Read elements per lane and program width, not the constants.

Moore Threads S5000 is a two-case rule, and a toolchain floor

Tiling pays at 64 heads only. Against a 1332 GB/s ceiling at 32768 tokens the
generic kernel reaches 87.4% at 64 heads but 93.6% at 128, and a full
TPP x num_warps sweep at 128 heads finds nothing above 0.94x — there is no
headroom left there. The tests and benchmark use 64 and 128 heads only, so this
is an exhaustive two-case rule rather than a fitted threshold. At 64 heads the
best cell is TPP=4/num_warps=4, 91.6% of ceiling.

The 64-head dip is real and was checked before being built on: 48 heads 90.8%,
64 heads 87.4%, 96 heads 90.8%, 128 heads 93.6%, with 0.1-0.6% spread over
three round-robin rounds and -0.3% ceiling drift. It is not monotonic, so "the KV
straggler is amortised as head count rises" is wrong, and it is not a
power-of-two stride effect either — 64 is the worst and 128 the best. The
mechanism is unexplained; it is also not load-bearing, since only 64 and 128
occur.

The token threshold is 512, chosen from where the measurement stops being noise:
below 192 tokens three round-robin repetitions spread 12-34%, so the apparent
0.82x at 64 tokens and 1.02x at 96 are both unusable. From 192 up the spread falls
to <= 3%, and the crossover sits between 256 (1.00x) and 512 (1.03x).

This override requires flagtree >= 0.6.1+mthreads3.6, and that is a real
floor, not a precaution. On older builds the tiled kernel does not compile in
any configuration — five were tried, including TPP=1/num_warps=1, whose 16
elements per lane match the generic launch that works fine, and all five abort
with llc code -6. The defect is an instruction-selection failure in the llc
shipped with MUSA toolkit 4.3.5 (MTGPU DAG->DAG Pattern Instruction Selection,
SelectionDAGISel::CannotYetSelect on a v8bf16 LSU_LD_CACHE_HINT), reached
because the older FlagTree wheels ship no bin/ and Triton falls back to the
system binary. 0.6.1+mthreads3.6 bundles a working llc at
triton/backends/mthreads/bin/llc, md5 cec9ff66714e311670b9412ec760e4aa —
byte-identical to the build the kernel was validated against, which is how the
version floor was confirmed rather than assumed. The two binaries both
self-report LLVM 14.0.0, so a version string cannot tell them apart; compare md5.
On a build without it the dispatch simply routes to the generic kernel, so
nothing here is reached and nothing breaks.

This repository pins no compiler — there is no backends.yaml and no flagtree
dependency — so which Triton runs is the installation's choice. That is also why
mthreads is added to the fp8 vendor whitelist here: I argued against it twice on
the grounds that opening the gate would turn a clean skip into compile failures
on the pinned FlagTree, which was importing a constraint from FlagGems that this
repo does not have.

Measured on the card with 0.6.1+mthreads3.6, two independent benchmark runs
agreeing within 0.5% on every shape:

tokens	64 heads (generic → tiled)	128 heads (control)
1024	0.965 → 1.040	1.074 → 1.064
2048	1.115 → 1.088	1.077 → 1.078
8192	1.144 → 1.205	1.088 → 1.087
32768	1.197 → 1.252	1.124 → 1.122
65536	1.208 → 1.253	1.151 → 1.150
98304	1.182 → 1.222	1.114 → 1.121
131072	1.200 → 1.264	1.138 → 1.135

The 128-head column is a built-in control and it holds: it takes the generic
path either way and is unchanged cell by cell, so the 64-head movement is the
tiling and not environment drift. Large shapes gain 3.4-5.3%, consistent with the
sweep's TPP=4/warps=4 figure of 1.06x against generic. The 2048 cell goes the
other way by 2.5%, which is inside this card's noise band but is not hidden here.

mthreads is added to the fp8 vendor whitelist, so the suite now runs on S5000
instead of skipping 108 tests. That gate is what hid the hardcoded device = "cuda" underneath it: nothing on that card ever reached the operator, so neither
defect was visible.

Accuracy on the card, excluding the four largest token counts: 76 passed, 16
skipped, 0 failed — every 64-head case at 512 tokens and above, which is the
set that takes the tiled path, so the new llc produces a kernel that is not only
buildable but numerically right. What happens at the four excluded shapes is in
the last section, under the memory-guard fix; it is not a property of this
override.

T-Head PPU-ZW810E: the operator could not be built at all

Not a tuning problem. tl.float8e4nv is refused for this card, so the generic
kernel does not compile — there is no "before" to compare against.

The card reports capability (8, 0) and its Triton target is ppu:80, and the
fp8 policy is inherited from the NVIDIA backend at two levels. FlagTree's
third_party/ppu/backend/compiler.py adds fp8e4nv to supported_fp8_dtypes
only when capability >= 89; that option is overridable, and overriding it moves
the failure into the PPU's own lowering pass, which fails with
`LLVM ERROR: Conversion from/to f8e4m3nv is only supported on compute capability

= 89— an NVIDIA threshold, and an NVIDIA concept, enforced insideConvertTritonGPUToLLVMPPU`. So the format is refused by inherited policy rather
than shown to be unsupported.

The override changes one thing: the seven unrolled quantisation groups call
_f32_to_e4m3_bits instead of .to(tl.float8e4nv). Everything else is the
generic kernel. The encoder is integer arithmetic — which is what a software fp8
conversion does anyway below sm_89; torch's own cast on this card takes such a
path. It is bit-identical to torch.float8_e4m3fn over 65536 values, half of
them drawn from the subnormal range below 2^-6, which is the part a naive encoder
gets wrong: a first version flushed subnormals to zero and also mis-saturated
exponent 15 with mantissa 0, which is the legal value 256 rather than an overflow.

Accuracy suite on the card: 92 passed, 16 skipped (the 16 are the vLLM
cross-check; the vLLM installed there is 0.19.0, which predates this operator).
That run predates the memory-guard removal described in the last section, so on
this card as on the others, whether the four largest shapes fit is a separate
question from whether the override is correct.

Performance is not measurably hurt by the software conversion — unsurprising for a
memory-bound operator, but measured rather than assumed: 96.0% of a warm
2097 GB/s ceiling at 32768 and 131072 tokens, ramping from 63.9% at 1024. Small
shapes need a cache flush between iterations to measure at all: this part has
roughly 32-64 MiB of cache (a 512 MiB copy reads 2097 GB/s, a 32 MiB copy 7140),
and the operator rewrites q in place, so at 1024 tokens its 67 MB footprint sits
right at the boundary and an unflushed measurement reads high.

The ZW810E baseline is T-Head's own vLLM port, used unmodified. That card's
pip index serves only +ppu* builds, and the newest,
vllm==0.20.1+v0.1.0.ppu2.1.0, registers
_C::fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert — a vendor C++
implementation, which does its FP8 conversion outside Triton and is therefore
untouched by the fp8e4nv gate that blocks the Triton path on the same card. It
launches, and it agrees with this operator under the accuracy suite's own
k_cache_compare at 64x64, 777x64 and 4096x128.

This is the only one of the baselines that needs no local patch at all: with
the vendor wheel on PYTHONPATH the benchmark resolves torch.ops._C
unmodified. Figures are from two independent runs agreeing within ~1% on every one
of the 22 shapes, with a known-answer device check run clean before and after each
— that check is there because this card intermittently enters a state where it
returns wrong answers, cause unidentified, so every measurement on it is
repeated rather than trusted once.

The box also carries a stock vllm 0.19.0+cu130, which predates the operator; an
earlier revision of this description reported ZW810E against vLLM's portable XPU
Triton implementation carrying our own E4M3 encoder, because that stock build was
mistaken for the only vLLM available. Those numbers are withdrawn in favour of
these.

Delete _thead/fused/ once the PPU backend either lowers fp8e4nv or reports a
target that permits it. Nothing else in it differs from the generic
implementation.

Ascend 910B: eight compiler limits, then 6.7x / 9.9x

Also an enablement before it is an optimisation, and for more reasons than one.
Four separate constructs the generic kernel relies on are unavailable here, and
they are the same four that stop vLLM's portable Triton fallback from
compiling on this card — taken verbatim from upstream and tried at this
operator's own shapes:

x_clamped.to(tl.float8e4nv)
    -> 'hivm.hir.vcast' op currently don't support cast
       float_to_UNKNOWN_rintmode / unrecognized float type: 'f8E4M3FN'
token_bf16_ptr.to(tl.pointer_type(tl.bfloat16))
    -> Casting pointers with unmatched bitwidth!

with encoded_scale.to(tl.uint8) from a float and an inexact tl.exp2 waiting
behind them. So 910B has no baseline of any kind — no vendor kernel and no
portable fallback — which is a stronger statement than the "no vendor op" that
applies to BW1000 and S5000, and is exactly why this file exists.

The override therefore carries: an integer E4M3 encoder (same reason as ZW810E),
the UE8M0 scale built from exponent bits rather than a float-to-int, a 2-D offset
for the RoPE pairs instead of a pointer cast, and a launch issued in chunks of
65535 programs because the runtime refuses anything wider.

Tuned from 26.8 GB/s — the first version that ran — in four steps, each one a
measurement correcting a guess:

step	what	64 heads	128 heads
—	one (token, head) unit per program	26.8	26.8
1	32-unit flat Q tile	56.9	70.5
2	E4M3 subnormals rounded arithmetically	81.6	86.8
3	Q tiled over heads within a token	114.6	186.4
4	KV quantised as one [8,64] tile, 3 dead branches cut	179.2	265.1

6.7x and 9.9x, confirmed two ways: a private loop timing both versions
alternately on the same tensors, and the repo's own Benchmark.run() with the
old operator bound into the empty torch_op slot. The two agree at every shape
from 1024 tokens up.

Two of those steps are worth stating because the mechanism was not the obvious
one. Step 2 looked like an encoder rewrite; the actual cost was three shifts by
a per-lane amount in the subnormal path, which this vector unit has no
instruction for and so scalarises — an E4M3 subnormal is m * 2^-9, so adding
2^23 and subtracting it again forces the same rounding with no shift at all, and
quantisation went 1.276 → 0.627 us/token. Step 3 looked like a tiling win and was
really a gather removed: every head of a token shares one position, but a flat
(token, head) tile makes position_id a vector, so each unit re-gathered the same
256 bytes of cos/sin once per head. Tiling heads of one token makes it scalar,
11.43 → 4.15 ns/unit at 64 heads and 16.81 → 3.99 at 128 — which also explains
what had looked unexplainable, the same kernel being 48% slower at 128 heads than
at 64.

Then the Q and KV halves were merged into one launch branching on the program
id, rather than two. A launch costs ~450 us on this card, so this is 1.94x below
~256 tokens and 1.02x at 1024 and above, bit-exact including mixed negative
slots. It also stops num_tokens being a constexpr, so shapes no longer force
recompiles — the suite runs in 5.15s against 13.53s.

That merge exists because of a lesson about A/B measurement, not because of a
kernel insight: the first before/after sampled only the shapes the tiling
targeted, all ≥1024 tokens. Run against the benchmark's own shape list, the tuned
version was ~1.9x slower below 64 tokens — two launches where the untuned one
had issued a single fused launch. Those are the decode-side shapes. A before/after
that samples only the shapes the optimisation targeted will not show what it cost
elsewhere.

The baseline here is an eager torch_npu composition, since neither of the
usual two exists: npu_rms_norm for the weightless RMSNorm (not a naive five-op
version), elementwise RoPE, the same integer FP8 encoder the kernel uses
(torch_npu cannot cast to float8_e4m3fn at all, so there is no alternative),
and an advanced-indexing scatter into the paged cache. It is validated against
the test file's own oracle with the operator taken out of the middle, so the
argument is not circular — k_cache bit-identical on every shape checked, mixed
negative slots included. It settles at ~11x on large shapes; the 7.6-9.9x at small
shapes is dozens of framework launches against one.

It cannot run 98304x128 or 131072x128, and not because of any single
allocation: the composition holds q.float(), npu_rms_norm's fp32 output, the
fp32 rope intermediate and the bf16 write-back live at the same time, about six
times q's own size, so the peak is 75 and 100 GiB against 60.96 available. Those
two are skipped by predicting the footprint, not by catching the failure — on
this backend an allocation error does not reliably land on the call that caused
it, and one left in flight poisons the rest of the process, which is how a stray
aicpu timeout corrupted an earlier sweep of this same baseline.

Two of 22 shapes, and not the large end. 131072 tokens at 64 heads runs and
is reported; only its 128-head partner is missing, which is why the 128-head
column above is labelled to 65536 rather than averaging over the survivors and
calling the range 32768-131072. The operator does both missing shapes in 97 and
129 ms, so the card is not the limit here — the composition is.

Two ways to make them fit were considered. Chunking over tokens multiplies the
framework launch count by the chunk count; the overhead looks small at these
sizes but it moves the ratio in our favour, and it was not measured, so it is not
relied on. Feeding npu_rms_norm bfloat16 would fit in 36 GiB and moves the
ratio against us, which is the safe direction, but it rounds the normalised
value before the rotation and stops matching the kernel's arithmetic. Neither is
needed to support a claim made here, so the two cells stay empty and say why.

ascend is in the fp8 vendor whitelist for the same reason as mthreads and
metax: the card reports capability (0, 0), so the NVIDIA rule skips the whole
file.

Two things this card taught that are worth a reviewer's time. First, torch_npu
cannot cast to float8_e4m3fn, so the test's torch oracle needs a CPU fallback
for that one step — one of the seven fixes below. Second, twice during this work
the operator was right and the oracle was wrong (exp2, rsqrt); the only
thing that ever settled such a case was dumping intermediates and diffing
bit-for-bit against an exact reference until the first divergence appeared.

Everything else tried on this card, and rejected on measurement: packing four
FP8 bytes into an int32 store (0.19x — five times slower, and bit-exact, so
the packing was right); num_warps 2 or 4 (identical to three decimals, here as
on every kernel in this PR); tiling KV over tokens as well (1.05-1.06x, ~3% of the
operator, against having to solve the scattered cache store); a 16-bit encoder
(0.99x, after a synthetic microbenchmark promised 1.33x — its op mix was not
the encoder's); capping the grid at the vector core count and looping inside
(nothing, and 20/32/40/48/64/128/256 were all measured); an exact reciprocal
multiply instead of the divide (nothing, within 1%); and TLE
dsa.alloc/dsa.copy staging the Q tile through UB (1.00x / 0.99x,
bit-exact — on this card the vector unit can only reach UB, so tl.load already
stages through it and the explicit copy restates what the compiler does anyway).
Nothing on that list is untried.

Dispatch thresholds

Below a threshold the generic kernel wins, because a wider block costs more to
launch and TPP masks off most of every program when there are fewer tokens than
that to fill it. The three tiling cards want three different rules and all three
were measured:

• C550: num_tokens < 512 falls through, on both head counts (256 tokens:
  0.96x/0.99x; 512: 1.03x/1.04x). The launch floor moves ~26 us → ~32 us.
• BW1000: crossover at 256 tokens for 64 heads and 128 tokens for 128 heads —
  16640 and 16512 programs respectively, so the quantity the two agree on is
  num_tokens * (num_heads + 1), and the threshold is a program count, 16384.
  A flat token threshold would forfeit the 1.16x-1.32x available between 128 and
  256 tokens at 128 heads. The launch floor here is much higher in absolute terms:
  80.5 us → 90.2 us at one token.
• S5000: num_heads > 64 or num_tokens < 512 falls through — the two-case
  head rule above, plus the token threshold from where the measurement stops being
  noise.

Correctness

Each override is verified against the torch reference, not against the generic
kernel, and that distinction turned out to matter. On C550 the tiled and generic
kernels are bit-identical. On BW1000 they are bit-identical under torch 2.10.0 /
FlagTree 0.6.1a1 — 32 shapes, 10 repetitions each, zero differing bytes — but
under the older combination that card's pinned stack installs (torch 2.9.0 /
FlagTree 0.5.1), q differs on 5 of 16744448 elements by up to 2 bf16 ULP at 511
tokens and one FP8 cache byte differs at 1000 tokens. The same appears on 910B
between the pre- and post-tiling versions: 66 of 33.5M elements at 2-3 bf16 ULP,
at magnitudes near 1e-6.

That is expected once stated precisely: a 2-D [TPP, 512] tile assigns lanes
differently than a 1-D 512 block, so the RMSNorm variance rounds differently, and
how that lowers is the compiler's choice. Near zero, bf16 is sparse enough that
"3 ULP" is 1.2% of a value whose absolute error is 3e-8. So
test_backend_override_matches_reference compares each implementation against the
reference at the tolerances the other tests use, rather than requiring two Triton
implementations to agree more closely than either agrees with the oracle. It runs
on shapes the other tests do not reach — 511, 512, 777, 1000, 4096 tokens, chosen
for not being multiples of a plausible tiling — and skips where no override is
registered.

On BW1000 both suites are green: 60 passed / 0 failed / 0 skipped with the
override active, and the same with the dispatch threshold temporarily set to 0,
which is the only way to exercise the tok_ok masking at 1, 4 and 17 tokens,
where 7 or more of each program's 8 slots are empty. At the production threshold
those shapes dispatch to the generic kernel and the masking is never tested.

On S5000, each of the four dispatch branches was checked directly: 777x64 and
32768x64 take the tiled path and match the generic kernel byte-for-byte in the FP8
cache with q within one bf16 ULP, while 511x64 (below the token threshold) and
777x128 (above the head threshold) are bit-identical, which is what confirms
the fallback really calls the generic implementation.

Verification

Every figure and pass/fail above was produced by running this branch's code on the
physical cards, not inferred. Where a measurement is not from this repo's
benchmark — the 910B rows, the BW1000 generic-vs-override bandwidth pair — that is
said at the point it appears.

select_backends.py picks non-NVIDIA backends from the changed paths under
src/flaggems_vllm/runtime/backend/_<vendor>/, so this PR should select metax,
hygon, mthreads, thead and ascend on its own, provided each is marked enabled in
the CI registry — I cannot see .ci/flaggems-registry/.github/backends.json from
here to confirm which are. The run is currently action_required because this
comes from a fork, so it needs a maintainer to approve it before anything
executes. Happy to attach raw logs from the cards in the meantime.

The seven test and benchmark fixes

A vendor override of this operator was unreachable from its own test and
benchmark. SpecOpRegistrar writes vendor implementations into
flaggems_vllm's module globals, but this operator's test called
flaggems_vllm.ops.<op> and its benchmark did set_gems(flaggems_vllm.ops.<op>) —
attributes of the ops submodule, which the registrar never touches. So both
always ran the generic kernel, silently, and no override for this operator could
ever have been measured or validated. Both now use the top-level binding
(flaggems_vllm.ops_experts_impl, flaggems_vllm.weight_norm,
flaggems_vllm.sparse_attn_triton, flaggems_vllm.w8a8_block_fp8_matmul are the
existing examples). Of the 29 set_gems call sites in benchmark/, 28 use the
top-level binding; the only other .ops. one is
test_flash_attn_varlen_opt_init_func.py, whose operator has no vendor override,
so it is unaffected. This is a prerequisite for the five overrides above to have
any effect.

device = "cuda" was hardcoded in four places. MACA aliases cuda, so this
is invisible on MetaX; on MUSA every case dies at its first allocation with
NotImplementedError: Could not run 'aten::empty.memory_format' with arguments from the 'CUDA' backend, before reaching the operator at all — 60 of 92 tests,
none of them for a reason related to the operator. Now flaggems_vllm.device,
which 60 of the 69 test files already use. Same family as the capability gate
below: both silently exclude non-NVIDIA backends, and neither can be noticed from
an NVIDIA machine.

The correctness oracle depended on test collection order. ref_impl ran
vLLM's C++ kernel instead of this file's torch reference whenever
torch.ops._C happened to carry the op. Nothing in this module imports vLLM — but
test_cp_gather_indexer_k_quant_cache.py and test_cutlass_scaled_mm.py do so at
module level and both sort earlier, and pytest imports every collected module
before running any test. So the oracle was vLLM's kernel when the suite ran and
the torch reference when the file ran alone, with nothing reporting the switch.
Worse, the vLLM path cannot execute on MetaX C550, so simply adding the missing
import — the obvious fix — would have broken the suite there. ref_impl is now
always the torch reference, and the vLLM cross-check is its own skipif-gated
test, which is how test_persistent_topk and test_top_k_per_row_decode are
already structured.

is_support_fp8e4nv() gated on get_device_capability() >= (8, 9). That
threshold means "Ada or newer" on NVIDIA only; other vendors report their own
major/minor on a different scale — C550 reports (8, 0), S5000 (3, 1), ZW810E
(8, 0) and 910B (0, 0), and all four convert fp8e4nv correctly (two of them via
the integer encoder in this PR). The check now consults an explicit vendor list
before falling back to the NVIDIA capability rule, and stays local to these two
files.

The torch oracle's FP8 cast cannot run on every backend. torch_npu has no
cast to float8_e4m3fn at all, so on 910B the reference itself raises before the
operator is reached. That one step now falls back to CPU, and the oracle's
power-of-two scale is built exactly rather than through exp2, which is off by one
ULP for integer arguments there — in both tl and torch. Both are oracle fixes
and neither touches the operator.

The shape exclusion was hardcoded and silent. if (num_tokens == 98304 or num_tokens == 131072) and n_heads == 128: return — a bare return, so those four
of 60 cases were reported as passed, and having been added in the commit that
introduced the operator, they had never actually run on any card. The list is also
wrong in both directions on a 64 GiB part: it excludes the 128-head shapes and does
not exclude 64 heads x 131072, which OOMs there. Removed outright, with nothing
in its place: those cases now run, and on a card that cannot hold the reference
they fail visibly rather than reporting success.

Two softer treatments were tried on this branch and both are reverted. Evaluating
the reference and the comparison in token slices let the large shapes fit, but it
changed what the oracle computes — and an oracle is the wrong place to economise:
if the whole-tensor path is right and the kernel has a size-dependent fault, a
sliced reference agrees with it and hides the one bug worth finding. A memory
budget check that skipped shapes instead was closer to honest, but a shape that
cannot be verified on a given card should say so rather than disappear.

What that costs, stated plainly, because it is more than "those shapes go red".
On 910B it is 11 OOM cases out of 92, with the operator not on the stack — the
oracle's xf.pow(2).mean(...) builds a 32 GiB float32 copy of a 16 GiB q.
On S5000 it is worse than a per-case failure: MUSA errors are sticky, so the
first failure poisons the context and every later GPU call in the same process
raises too. The same command run twice gave 12 failed / 80 passed and 55 failed /
37 passed — the difference being only where in the collection order the first
failure landed. Exclude the four largest token counts and it is 76 passed / 16
skipped / 0 failed. The first failure there is itself a known vendor defect
(torch_musa#146): boolean
masked assignment fails past 2^31 elements, and it sits directly on
assert_close's diff-reporting path, so the failure surfaces as
RuntimeError: Comparing with the actual difference never printed. Neither of
those is caused by this change; both are made visible by it, which was the point.

The benchmark probed torch.ops._C without importing the library that
registers the op, and torch.ops._C gives no hint that nothing did. Where the
reference is installed the benchmark reported it missing and skipped, so no
comparison ran at all. It now does exactly what the three sibling benchmarks that
use torch.ops._C do — top_k_per_row_decode, persistent_topk and
cutlass_scaled_mm — a single import vllm._custom_ops under try/except, then
hasattr. On MetaX C550 that entry point is also the one that works: import vllm._C fails on libcudart.so.13, while vllm._custom_ops registers all 57 ops
through mcoplib.

That guard catches RuntimeError as well as ImportError/AttributeError, which
is a deviation from the three siblings and is deliberate: on a box with two vLLM
platform plugins registered, the import raises RuntimeError: Only one platform plugin can be activated, but got: ['fl', 'musa'], and with the narrower guard that
aborts collection of the entire file rather than degrading to "no baseline".
A misconfigured optional dependency should not be able to do that. It is not
widened to bare Exception, which would swallow real problems.

A reference can also be registered and still not run, so the benchmark
distinguishes the two. MetaX's build of this op returns mcErrorInvalidValue from
every launch on C550, and finding it that way would turn the benchmark red on the
vendor's defect, while skipping as "not installed" hides it. The reference is
wrapped so the first call reports the launch error as a skip reason. Making that
work took measuring how the failure propagates, and the result is worth recording:
on that backend a failed launch is not surfaced by synchronize(), nor by a
device-to-host copy, nor by a small allocation — it is surfaced by the next kernel
launch, which is why it otherwise lands on do_bench's 256 MB L2-flush
allocation, several frames from the cause.

Two things reviewers should know

Vendor discovery is a no-op where mcoplib is absent, so the benchmark skips
exactly as before on such machines. For anyone running it on a machine that does
have it: on C550 that kernel cannot launch (reported as
MetaX-MACA/mcoplib#59), and the
benchmark will now say so rather than silently comparing against something else.

The ~140-line tiled kernel now exists in three copies — _metax/fused/,
_hygon/fused/ and _mthreads/fused/ — with identical kernel bodies; only the
launch constants and the dispatch rule differ per card. (_thead and _ascend
are not copies of it: one is the generic kernel with a different FP8 conversion,
the other a differently-shaped kernel built around this card's compiler limits.)
Three copies will stay in sync by accident rather than by construction, and the
case for factoring it out is stronger than it was with two. The alternatives are a
shared module or one backend importing another's, and neither looked obviously
right in this layout, so raising it rather than deciding it unilaterally — happy to
factor it out whichever way you prefer.

[CLAassistant]

All committers have signed the CLA.