[KMCompiler][MetaX][Iluvatar][Hygon] Add linalg_lstsq MetaX, Iluvatar and Hygon backend support by cheersluvs · Pull Request #5390 · flagos-ai/FlagGems

cheersluvs · 2026-08-11T07:57:39Z

What

Enables linalg_lstsq on MetaX and Iluvatar. The generic operator,
added for NVIDIA in #4938 and extended to Ascend in #5173, is not modified —
every device-specific value is applied by rebinding attributes on
flag_gems.ops.linalg_lstsq, which works because the generic driver reads all
of them at call time.

file	change
`_metax/ops/linalg_lstsq.py`	new
`_iluvatar/ops/linalg_lstsq.py`	new
`_metax/ops/__init__.py`, `_iluvatar/ops/__init__.py`	+2 each (registration)
`tests/test_linalg_lstsq.py`	reference helper, see below
`ops/linalg_lstsq.py`	unchanged

The two vendors are together because they share the test-reference change and
the Iluvatar half depends on it; splitting them would mean a stacked PR.

Copying the ~250-line compact-WY driver into the vendor directory would have
been the alternative, and it would drift the moment the generic side changes.

Sizing every tile to the device

Upstream's block sizes assume NVIDIA-sized shared memory. A MetaX part has
64KB, and three separate tiles overflow or exactly fill it. All three are
now derived from the limit Triton actually enforces — its driver's
max_shared_mem, not torch's shared_memory_per_block, which reports the
49152 static default on NVIDIA while Triton opts in to 164–228KB, so deriving
from torch would needlessly shrink tiles on an H20.

constant	tile	upstream	derived
`_WY_BLOCK_C`	compact-WY update	73728 B (8KB over)	24576 B
`_TARGET_TILE_BYTES`	`_panel_qr_kernel`	65536 B at NC=33	32768 B
`_TARGET_STACK_ROWS`	`_reduce_qr_kernel`	65536 B at NC=33	32768 B

The budget is half the limit in each case, and that is the substance of
the change, not caution. A 65536 B tile passes the > check with nothing left
for launch overhead, and leaves room for only ONE block per SM. Measured on a
C550, BLOCK_C=16 (24576 B, two blocks resident) ran the WY update in 0.0797
ms against BLOCK_C=32 (40960 B, one block) at 0.0977.

The two 65536 B tiles were invisible on the hardware I had: they fit exactly,
so the suite passed while the tall path ran at one block per SM throughout.
The maca3720 runner, added mid-review, failed on them at (256,32) with
Required: 66560, Hardware limit: 65536 — and fixing it turned out to be
worth +19% fp32 geomean, not just a green check.

Configuration is derived on first call, never at import: querying the Triton
driver while flag_gems is still loading its vendor ops is not something a
vendor module should do.

Two more MetaX device facts

float64 tl.dot is miscompiled. Relative error ~1.0 at every operand shape
tested, including 16×16×16, with and without input_precision="ieee", while
the same contraction written as rank-1 updates is bit-exact (0.00e+00). fp32
tl.dot is correct (1.1e-07 with ieee). _wy_update_metax branches on the
COMPUTE tl.constexpr, so this is resolved at compile time and fp32 keeps
the tl.dot path. This is a MACA compiler bug, not a design choice — if it
is fixed, deleting the COMPUTE == tl.float64 branches restores the fast path.
I have a self-contained reproducer and am happy to share it.

4 KB/thread of private (register-spill) memory, a driver setting
(insmod metax.ko pri_mem_sz=…), which the monolithic reduce kernel exceeds at
BLOCK_NC=128. It surfaces as the misleading Triton Error [MACA]: memory size or pointer value too large to fit in 32 bit; the real message only appears in
captured stdout. num_warps=8 fits and is numerically correct (4.1e-07) but
costs 25× (14.9 ms vs 0.60), and num_warps=16 exceeds the device's 512
threads/block. So _TALL_MAX_NC_F32 is capped at 64 and those shapes route to
compact-WY, which has no size ceiling by design.

Iluvatar

One device fact, and it is a compiler limit rather than a numerical one:
float64 tl.dot does not compile on a BI-V150.

test_linalg_lstsq_tall_blocked_fp64
triton.compiler.errors.CompilationError: at 39:16:
    Wacc += tl.dot(tl.trans(Vb), Tb, input_precision="ieee")

Every other fp64 case passes, because the monolithic and blocked-TSQR paths
contain no tl.dot at all — only compact-WY contracts with one, and all three
of the operator's tl.dot calls live in that single kernel. So the port is one
kernel and one rebind: _wy_update with a float64 form expressing each
contraction as P rank-1 updates (P is the panel width, 16, not a problem
dimension), selected by the COMPUTE constexpr so float32 keeps tl.dot.

Nothing else is rebound. No block-size configuration is shipped for this
backend: none has been measured on the device, and unmeasured tuning does not
belong in an override.

float64 is skipped here, so that branch is not exercised in CI. The device
declares support_fp64 = False and means it — torch.matmul reports "gemm of
double is not supported on CoreX", cuSOLVER has neither Dormqr nor Dorgqr,
every float64 copy warns "limited support", and a 256×256 float64 solve returns
NaN even with this kernel in place. _require_dtype now consults that flag
first (the old operator probe used a 4×2 solve, which routes to the monolithic
path and so answered for one of four paths — it said "supported" while the
256×256 case produced NaN). The branch is kept because the compile error it
removes is real, it costs float32 nothing (COMPUTE is a tl.constexpr), and
it is what makes compact-WY work if this device gains usable float64.

The rank-1 kernel is near-identical to the MetaX one, which is duplication I
would rather not have. Sharing it would mean either a cross-vendor import (the
MetaX module rebinds its own shared-memory tuning at import, so importing it
would apply MetaX block sizes to Iluvatar) or a USE_DOT: tl.constexpr on the
generic kernel. If a third vendor hits the same gap, the generic flag is the
right answer.

The one change outside the vendor directories

_ref_and_gems now computes its reference with a float64 solve on the CPU.

Where to_reference keeps the tensor on the device, torch.linalg.lstsq in
the test is the vendor's kernel, which is neither reliably accurate nor
reliably present. Against a float64 CPU solve at κ ≈ 1.6–2.0:

	max abs error
FlagGems	5.2e-08
MACA `torch.linalg.lstsq` (fp32)	1.3e-04

Two cases were failing on the reference's error rather than ours. Upcasting
alone is not enough — on maca3810 the fp64 device solve does not exist and
raises invalid device function — so the solve moves to the CPU outright. The
result is placed where gems_assert_close expects it (CPU under --ref=cpu,
device otherwise) and its dtype is gated on fp64 support, so devices without it
(Ascend 910B) get an fp32 reference exactly as before.

Four more tests hand-rolled the same pattern and bypassed that helper. On
Iluvatar square_wy_fp64 failed inside its own reference with
cusolver error ... cusolverDnDormqr_bufferSize, because that shim has no
float64 QR — nothing an operator override can reach. They now use the same CPU
solve. Nothing is skipped and no assertion is dropped; near_singular keeps
its reference in the input dtype because it asserts the result's dtype matches
torch's.

Results

pytest tests/test_linalg_lstsq.py:

device	result
MetaX C550	76 passed, 0 skipped
Iluvatar BI-V150	71 passed, 5 skipped (4 float64 + 1 complex), was 74 passed / 2 failed

MetaX CI runners maca3720 and maca3810 green; python-op green.

MetaX --mode kernel, 23 shapes per dtype:

	fp32	fp64
geomean speedup	8.589×	2.421×
arithmetic mean	20.619×	5.916×
at or above 1×	22/23	15/23

Iluvatar --mode kernel, float32: geomean 5.226×, arithmetic 16.444×,
18/23 at or above 1×, range 0.599–79.867. Every sub-1× row there is a
compact-WY shape, which is the untuned path on this backend — deriving its
block sizes from the device, as the MetaX half does, is the obvious follow-up
and is deliberately not attempted here without measurements.

Where MetaX loses:

fp64, large NC (1024², 2048², 4096×512, 512×1024, 1024×2048) route
through compact-WY, which is exactly where the rank-1 rewrite replaces
tl.dot. This is the measured price of the fp64 tl.dot bug.
small shapes where torch is already sub-millisecond — launch overhead,
unrelated to that workaround.

Note for reviewers

_metax/ops/linalg_lstsq.py raises a RuntimeError if fp32 and fp64 derive
different BLOCK_C. The generic driver computes the WY update grid from the
module constant _WY_BLOCK_C while taking BLOCK_C from _wy_cfg; those
agree upstream, so the mismatch is dormant there, and rebinding both keeps them
agreeing here — but only while both dtypes derive the same value (they do on
C550: both 16). On a device where they diverge, the right fix is a one-line
generic change making the grid use the per-dtype value; failing loudly beats
silently under-updating the trailing block.

Adds a MetaX override for linalg_lstsq. Every device-specific value is applied by rebinding attributes on flag_gems.ops.linalg_lstsq at import, so the generic operator is not modified at all -- the driver reads each of them at call time, and duplicating its ~250-line compact-WY driver would only drift the moment the generic side changes. Three device facts, each measured on a C550 rather than inferred: 1. 64KB of shared memory per block. The tuned WY update config asks for 73728 bytes in both dtypes, 8KB over, and 19 of the suite's cases died with OutOfResources. BLOCK_C is now derived from the limit Triton enforces (its driver's max_shared_mem, not torch's static figure, which would make an H20 look like a 48KB part). The target is HALF the limit so two blocks stay resident per SM: bc=16 ran the update in 0.0797 ms against bc=32 at 0.0977, because only one block fits at 40960 B. 2. float64 tl.dot is miscompiled. Relative error ~1.0 at every operand shape including 16x16x16, with and without input_precision="ieee", while the same contraction written as rank-1 updates is bit-exact. _wy_update_metax branches on the COMPUTE constexpr, so fp32 keeps the tl.dot path unchanged -- it is correct here (1.1e-07) and faster. 3. 4KB/thread of private (register-spill) memory, a driver setting, which the monolithic path's reduce kernel exceeds at BLOCK_NC=128 -- surfacing as the misleading "memory size or pointer value too large to fit in 32 bit". num_warps=8 fits and is correct but costs 25x (14.9 vs 0.60 ms), and num_warps=16 exceeds threads-per-block, so _TALL_MAX_NC_F32 is capped at 64 and those shapes route to compact-WY, which has no ceiling by design. The one change outside the vendor directory is in the test helper: the reference is now built with to_reference(..., upcast=True). On backends where to_reference keeps the tensor on the device, torch.linalg.lstsq is the VENDOR's kernel rather than truth, and against a float64 CPU solve this op was off by 5.2e-08 where MACA's own fp32 lstsq was off by 1.3e-04 -- so two cases were failing on the reference's error, not ours. Upcasting is gated on fp64 support, so Ascend 910B keeps the fp32 reference. tests/test_linalg_lstsq.py: 76 passed on a C550.

`to_reference(..., upcast=True)` fixed the reference's precision but left the solve running on the vendor's device, which is the thing that was wrong with it. On maca3810 there is no fp64 `torch.linalg.lstsq`, so the reference itself raised `CUDA error: invalid device function` before the operator was ever called -- the failing case, (16,4) fp32 with NC=5, takes the monolithic path that this backend override does not alter at all. Compute the reference on the CPU in float64 and copy the result back to the device. That is the same reference on every backend, independent of vendor kernel accuracy and availability. The dtype of the copy is gated on fp64 support, so devices without it (Ascend 910B) get an fp32 reference exactly as before. Verified on a MetaX C550: 76 passed.

test-op.sh runs a second pass with `--ref=cpu`, which sets TO_CPU. In that mode gems_assert_close moves `res` to the CPU and asserts the reference is ALREADY there, so unconditionally copying the new CPU reference back to the device broke it -- the values matched exactly, only the device did not. Place the reference where the comparison expects it: CPU under TO_CPU, on the device otherwise. Verified on a MetaX C550, both passes green: the default run and `--ref=cpu --quick`.

The override queried Triton's driver for the shared-memory limit at import time, i.e. while flag_gems was still loading its vendor ops and before torch had touched the device. On the maca3810 CI runner, which uses the flagtree Triton backend, the next launch -- an ordinary `torch.randn` on the first line of the first test -- then failed with mcErrorInvalidDeviceFunction. Other metax PRs run that same call and pass, and a C550 with a different Triton never reproduced it. Nothing in this module needs to run before the operator is first called, so derive BLOCK_C lazily and rebind _WY_BLOCK_C there too. Import is now pure Python: no device access at all. The fp32/fp64 consistency check moves inside the same lazy path. Verified on a MetaX C550: 76 passed.

BLOCK_C was derived from the device's shared-memory limit for the compact-WY path, but the monolithic path's reduce tile never was. It stacks G = max(2, _TARGET_STACK_ROWS // NC) R factors into next_pow2(G*NC) x next_pow2(NC), which at the upstream 256 is exactly 65536 B for NC=33 -- the whole of a MetaX limit, with nothing left for Triton's ~1KB of launch overhead. A C550 fit it by sitting precisely on the boundary; maca3720 asks for 66560 and raises OutOfResources on (256,32). Derive _TARGET_STACK_ROWS the same way as BLOCK_C, budgeting half the limit (256 -> 128, tile 65536 -> 32768 B). Only the reduction fan-in changes, not the result. Measured on a C550, 76 passed, and the fp32 geomean is unchanged at 7.314 vs 7.241 (22/23 at or above 1x, was 21/23): the smaller tile costs some wide shapes and pays for itself on tall ones. fp64 never takes this path and is unmoved at 2.480.

_TARGET_STACK_ROWS covered the reduce tile, but the kernel maca3720 actually died in is _panel_qr_kernel, whose tile block_m x next_pow2(NC) is bounded by _TARGET_TILE_BYTES. Upstream's 96KB assumes NVIDIA-sized shared memory; on a 64KB part _choose_block_m yields 256x64x4 = 65536 B at NC=33, the whole limit, and the launch then asks for 66560. Four benchmark shapes -- (256,32), (4096,8), (1024,16), (2048,16) -- were sitting at exactly 65536. Budget half the limit, as the other two tiles do. This is a throughput fix as much as a correctness one: a 65536 B tile leaves room for ONE block per SM, and halving it puts two resident. Measured on a C550, fp32 geomean 7.241 -> 8.589 and arithmetic 17.326 -> 20.619, with (1024,16) +220%, (16,8192,16) +201% and (64,2048,16) +126%. Wide and underdetermined shapes give back up to 20%. fp64 does not take this path and is unchanged at 2.421.

float64 `tl.dot` does not compile on a BI-V150: test_linalg_lstsq_tall_blocked_fp64 triton.compiler.errors.CompilationError: at 39:16: Wacc += tl.dot(tl.trans(Vb), Tb, input_precision="ieee") Every other fp64 case passes, because the monolithic and blocked-TSQR paths contain no `tl.dot` at all -- only compact-WY contracts with one, and all three of the operator's `tl.dot` calls live in that single kernel. So the port is one kernel: `_wy_update` with a float64 form expressing each contraction as P rank-1 updates, selected by the COMPUTE constexpr so float32 keeps `tl.dot`. Nothing else is rebound. Unlike MetaX there is no shared-memory pressure here, and configuration not measured on the device does not belong in an override. Three tests also built their reference by calling `torch.linalg.lstsq` on DEVICE tensors, which runs the vendor's kernel rather than a trusted one. Iluvatar's cuSOLVER shim has no float64 QR, so `square_wy_fp64` failed inside its own reference with `cusolverDnDormqr_bufferSize` -- nothing an operator override can reach. They now solve on the CPU like `_ref_and_gems`, so the reference is the same on every backend. Nothing is skipped and no assertion is dropped. `near_singular` deliberately keeps torch on the device: it asserts gems returns the same four-tuple contract torch does -- solution, residuals, rank and singular_values -- so the device's own torch is the correct reference there. Measured on a BI-V150: 74 passed / 2 failed before, 75 passed / 1 failed with the override alone.

Both are harness portability, not operator behaviour, and both surfaced only once the suite ran on more than one vendor. square_wy_fp64 compares res against the reference with plain arithmetic rather than through gems_assert_close, so under `--ref=cpu` -- where _cpu_ref correctly leaves the reference on the CPU -- it mixed a cuda tensor with a cpu one. Route it through utils.to_cpu, which is what that helper is for. complex_fallback called torch.allclose on complex64 DEVICE tensors, which needs an elementwise complex abs. Iluvatar's runtime compiler cannot build one: `[IXRTC] nvrtcCompileProgram failed ... abs_kernel<std::complex<float>>`. The comparison gains nothing from running on the device, so keep the reference on the CPU and move the result to it.

`_require_dtype` asked only the operator, via a 4x2 probe. That probe routes to the monolithic path, so it answers for one of four code paths and cannot speak for the others: on an Iluvatar BI-V150 it succeeds while a 256x256 float64 solve returns NaN. The device is unambiguous that float64 is not available: * flag_gems.runtime.device.support_fp64 is False * float64 tl.dot does not compile * torch.matmul: "gemm of double is not supported on CoreX" * cuSOLVER has neither Dormqr nor Dorgqr * every float64 copy warns "Limited support for torch.double" So consult the declared flag first and skip float64 there, matching the suite-wide convention (ALL_FLOAT_DTYPES already gates on fp64_is_supported). The operator probe stays as a second gate, for backends that do have float64 but lack this particular kernel -- Ascend, which answers NotImplementedError. The _wy_update override keeps its float64 branch: the compile error it removes is real, it costs float32 nothing (COMPUTE is a constexpr), and it is what makes compact-WY work if this device gains usable float64. Its module now says plainly that CI does not exercise that branch.

Hygon has 64KB of shared memory per block, and upstream's compact-WY update config asks for 2*BLOCK_R*BLOCK_C*esize + BLOCK_R*P*esize = 2*128*64*4 + 128*16*4 = 73728 bytes. 14 of the suite's cases died with triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 73728, Hardware limit: 65536 Every one is a compact-WY shape (square, underdetermined, tall blocked, rank-deficient square); the monolithic and blocked-TSQR paths fit and pass. Derive BLOCK_C from the limit Triton actually enforces -- 65536 here, giving 16 (24576 bytes), the same value and the same reasoning as the MetaX port. _wy_cfg is rebound as well as _WY_BLOCK_C, and that is load-bearing rather than defensive: _wy_cfg returns a LITERAL 64 for float64 while float32 reads the module constant, so rebinding the constant alone would leave float64 launching a 64-wide tile against a grid sized for 16, silently under-updating the trailing block. float64 passes at 64 on this device today, so that would have been a live regression, not a theoretical one. Nothing else is rebound. _TARGET_TILE_BYTES and _TARGET_STACK_ROWS sit at exactly the limit here, which is the zero-headroom that later broke a MetaX runner, but every shape using them passes on this device and a backend override is not the place for configuration measured somewhere else.

The shared-memory fix cleared all 14 OutOfResources cases -- CI reached 49 passed -- and then hit a different wall: triton/language/semantic.py:1445 AssertionError: Unsupported lhs dtype fp64 CompilationError: at 39:16 (_wy_update's tl.dot) That is Triton's own frontend allow-list, not the silicon: a local Hygon box with a different Triton build runs the same float64 tests through `tl.dot` and passes all 76. The CI runner also declares float64 supported, so the dtype gate does not skip it, and _gems_supports probes with a 4x2 solve that routes to the monolithic path and never reaches a `tl.dot` -- so nothing upstream of the kernel can see this. Carry the same float64 rank-1 form the Iluvatar override uses, branching on the COMPUTE constexpr so float32 keeps `tl.dot`. The kernel is byte-identical to that one, which is duplication worth naming: three backends now carry it, and a single `USE_DOT` constexpr on the generic kernel would replace all three copies with a one-line rebind each. That is the right shape once a fourth appears. Note this makes float64 use rank-1 on Hygon boxes whose Triton would have accepted `tl.dot`, trading some float64 throughput for one file that works on both Triton builds.

cheersluvs requested review from 0x45f, 103yiran, Caeruleann, bin913, douxetpur, hellojack163, huangyiqun, tengqm and w1120029931-bit as code owners August 11, 2026 07:57

github-actions Bot added tests vendor/MetaX size/Medium labels Aug 11, 2026

tengqm and others added 3 commits August 11, 2026 17:00

Merge branch 'master' into pr/linalg-lstsq-metax2

4251708

Merge branch 'master' into pr/linalg-lstsq-metax2

c90949b

Merge branch 'master' into pr/linalg-lstsq-metax2

0132e35

cheersluvs force-pushed the pr/linalg-lstsq-metax2 branch from fe90006 to 2669b66 Compare August 12, 2026 00:55

cheersluvs force-pushed the pr/linalg-lstsq-metax2 branch from 2669b66 to ccd5cf9 Compare August 12, 2026 00:56

cheersluvs added 2 commits August 12, 2026 09:07

cheersluvs marked this pull request as draft August 12, 2026 02:05

cheersluvs added 3 commits August 12, 2026 14:03

Merge branch 'master' into pr/linalg-lstsq-metax2

0c95075

cheersluvs closed this Aug 12, 2026

cheersluvs reopened this Aug 12, 2026

cheersluvs changed the title ~~[KMCompiler][MetaX] Add linalg_lstsq Metax backend support~~ [KMCompiler][MetaX][Iluvatar] Add linalg_lstsq MetaX and Iluvatar backend support Aug 12, 2026

github-actions Bot added the vendor/Iluvatar label Aug 12, 2026

cheersluvs force-pushed the pr/linalg-lstsq-metax2 branch from bae0a87 to b3195a2 Compare August 12, 2026 08:40

github-actions Bot added size/Large and removed size/Medium labels Aug 12, 2026

cheersluvs closed this Aug 12, 2026

cheersluvs reopened this Aug 12, 2026

cheersluvs added 2 commits August 12, 2026 17:32

github-actions Bot added the vendor/Hygon label Aug 15, 2026

cheersluvs changed the title ~~[KMCompiler][MetaX][Iluvatar] Add linalg_lstsq MetaX and Iluvatar backend support~~ [KMCompiler][MetaX][Iluvatar][Hygon] Add linalg_lstsq MetaX, Iluvatar and Hygon backend support Aug 15, 2026

cheersluvs added 2 commits August 15, 2026 15:26

Merge branch 'master' into pr/linalg-lstsq-metax2

d72e21f

cheersluvs marked this pull request as ready for review August 15, 2026 07:51

Navigation Menu

Uh oh!

[KMCompiler][MetaX][Iluvatar][Hygon] Add linalg_lstsq MetaX, Iluvatar and Hygon backend support - #5390

[KMCompiler][MetaX][Iluvatar][Hygon] Add linalg_lstsq MetaX, Iluvatar and Hygon backend support#5390
cheersluvs wants to merge 16 commits into
flagos-ai:masterfrom
cheersluvs:pr/linalg-lstsq-metax2

cheersluvs commented Aug 11, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

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

Uh oh!

What

Sizing every tile to the device

Two more MetaX device facts

Iluvatar

The one change outside the vendor directories

Results

Note for reviewers

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

cheersluvs commented Aug 11, 2026 •

edited

Loading