[KMCompiler][MetaX][Hygon] Token-tiled fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert, plus test and benchmark fixes

[cheersluvs]

Token-tiled backend overrides of
fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert for MetaX C550 and Hygon
BW1000, plus four fixes to the operator's test and benchmark that a card other
than an 80 GiB H800 exposes. Six files, all belonging to this operator — no shared
module added or changed.

Performance

Measured against achievable bandwidth on each card — a 512 MiB device-to-device
copy, measured rather than a spec figure:

	ceiling	generic	tiled
MetaX C550, 131072 x 64	1388 GB/s	1164 (83.9%)	1339 (96.5%)
MetaX C550, 32768 x 128	1388 GB/s	1224 (88.2%)	1354 (97.5%)
Hygon BW1000, 32768 x 64	1340 GB/s	604.8 (45.1%)	1182.9 (88.3%)
Hygon BW1000, 32768 x 128	1340 GB/s	609.2 (45.5%)	1201.4 (89.6%)

Through this repo's benchmark, --mode kernel --level core. On C550 the baseline
is MetaX's own compiled kernel (mcoplib 0.4.6); on BW1000 there is no vendor
implementation of this operator, so the baseline is vLLM's portable Triton
implementation of it, which is what vLLM itself falls back to on such hardware:

	64 heads	128 heads
C550 vs mcoplib, large shapes	0.872-0.888 → 1.004-1.016	0.945-0.974 → 1.047-1.053
BW1000 vs vLLM Triton, 8192-131072	2.02-2.04x	2.99-3.22x
BW1000 vs vLLM Triton, 1-17 tokens	1.20-1.54x	1.27-1.54x

Two caveats on the BW1000 column, because it is easy to read more into it than it
says. The 128-head figures are larger mostly because the baseline degrades
there — this operator holds 89.6% of achievable at both head counts while the
baseline falls from 45% to 29.6% — so the honest summary is "about 2x, more at 128
heads because the baseline drops". And the 1-17 token rows dispatch to the
generic kernel, so that gain is one fused launch beating the baseline's two and
has nothing to do with tiling.

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.

What the optimisation does

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=8 tokens of one slot, so every program is
uniformly all-Q or all-KV with no divergence, at num_warps=4: 16 elements per
lane instead of 8, in a 256-thread block instead of 64. Both parts have 64-lane
warps, so the same configuration is optimal on both.

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.

Dispatch thresholds

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

• C550: crossover at 512 tokens on both head counts (256 tokens: 0.96x/0.99x;
  512: 1.03x/1.04x), and 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.

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 combination backends.yaml pins for that card, which is what CI
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.

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. 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.

Both suites are green: 60 passed / 0 failed / 0 skipped on BW1000 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.

The four test and benchmark fixes

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. MetaX C550 reports (8, 0) and converts
fp8e4nv bit-identically to torch, yet the whole file silently skipped. The check
now consults an explicit vendor list before falling back to the NVIDIA capability
rule, and stays local to these two files — the same shape the four other files in
the repo already use. Those four carry the same bug; they are out of scope here.

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. Replaced with a budget measured
against free device memory and an explicit skip.

The reference and the comparison are evaluated in token slices. The reference
upcasts q to fp32, which at 131072 x 64 is a single 16 GiB allocation and the
actual OOM on a 64 GiB card; and assert_close allocates several temporaries the
size of its inputs — isclose alone needs a boolean mask over every element — so
at 131072 x 128 it asks for 12 GiB with ~54 GiB already held, and it is the
comparison that runs out of room rather than anything under test. Slicing along
tokens is exact here, because the operator is per-token independent. With the
temporaries bounded, the four excluded shapes fit and run: 60 passed, 0 skipped.

A second, independent reason to keep the oracle small: on one backend the
pure-torch reference returns wrong results for very large tensors — it disagreed
with the operator on 4.16% of elements at 3.22e9 elements evaluated whole, and
0.0002% evaluated in four slices, while the operator's own output was
byte-identical either way. An oracle has to be computed in a range where the
framework can be trusted, or the suite measures the framework rather than the
operator.

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 imports vllm._custom_ops — the entry point vLLM
itself uses, which vendor ports patch, so it resolves to the right library per
platform. On MetaX C550 that matters concretely: import vllm._C fails on
libcudart.so.13, while vllm._custom_ops registers all 57 ops through
mcoplib. A rejected import is logged rather than swallowed — a silent
except: pass is what makes a missing baseline indistinguishable from an
unimportable one.

A reference can also be registered and still not run, so the benchmark now
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 on the MetaX CI runner, which does not carry
mcoplib, so the benchmark skips exactly as before — verified on this PR's own
runs. For anyone running the benchmark 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 is duplicated between _metax/fused/ and
_hygon/fused/, identical in both. Both parts are 64-lane, so the two copies
will stay in sync by accident rather than by construction. The alternatives are a
shared module or one backend importing the other'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.

[cheersluvs]

Cross-reference for anyone landing here: this was the wrong repository. The operator lives in FlagGems-vllm, whose conf/operators.yaml declares it as the replacement for torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert — which is exactly the baseline this work was measured against.

Reopened there as flagos-ai/FlagGems-vllm#684, carrying the same two backend overrides and the same test/benchmark fixes, ported rather than copied: the package is flaggems_vllm, the generic op lives under ops/ instead of fused/, and that repo formats with black 24.8.0 at line-length 88.

The port turned up a fifth issue that does not exist here: there, this operator's test and benchmark called flaggems_vllm.ops.<op>, an attribute of the ops submodule, while the vendor registrar writes into the package globals — so no vendor override of this operator could ever have been reached by its own test or benchmark. Fixed in #684 as a prerequisite for the overrides to do anything.

One thing worth flagging before this is forgotten: the four test/benchmark defects this PR fixed are present in this repo too, byte-identical — the fp8 capability gate keyed on an NVIDIA-only (8, 9) threshold, the hardcoded H800 shape exclusion whose bare return reports 4 of 60 cases as passed without running them, the unchunked fp32 reference that OOMs on any 64 GiB card, and the torch.ops._C probe that never imports the library registering the op. They are real here and remain unfixed. Happy to open a separate PR for just those if that is useful.