[joshuaswarren]

The flash attention shader of llama.cpp (flash_attn_cm1) accumulates its P times V product in a half precision cooperative matrix, so its cmat_muladd has half A and B and a half accumulator. lower_matrix_instr in agx_nir_lower_cmat.c asserts that the accumulator is float or bfloat. On an M5 Max with the NIR validation abort fixed (#40), loading a model through llama.cpp's Vulkan backend then stops in that assertion in agx_nir_lower_cmat.c, in the cmat_muladd case.

The lowering already converts the accumulator to 32 bits for the hardware multiply-add and converts the result back to the result type. A half accumulator therefore needs no new lowering; this only allows it in the assertion.

Test: CmatLower.HalfInputsHalfAccumulator in agx_tests lowers a 16x16x16 half multiply-add with a half accumulator and checks that no cmat_muladd is left and that cmat_muladd_agx takes a 32-bit accumulator. A HalfInputsFloatAccumulator control passes before and after. Without the change the half case aborts with the same assertion text as on the device. Built with assertions on (b_ndebug=false): full agx_tests 505 passed, 1 skipped. A release build compiles the assertion out, so the test only fails on a build with assertions.

Not run: the llama.cpp shader on the device with this change (the device is not available right now), so later passes in the same shader are untested. The driver still advertises only float accumulation for half inputs; this PR does not change that.

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The change is limited to AGX cooperative-matrix shader compilation.

The PR appears safe to merge, with a non-blocking cleanup issue in the new test helper.

Findings

1.  Tests retain the type cache ▶

Summary

Allows half-precision accumulators through the cmat_muladd assertion. The existing lowering already converts them to float for the hardware operation and converts the result back.

• AGX lowers half-accumulator matrix multiply-adds.

_{Reviews (1) · Last reviewed commit: "agx: lower cooperative matrix multiply-a..." · Reviewed by Greptile}

[greptile-apps]

Tests retain the type cache

lowered_muladd_count() calls glsl_type_singleton_init_or_ref() on every run but never releases that reference. This keeps the shared type cache allocated after these tests finish, including across repeated runs. Call glsl_type_singleton_decref() after freeing the shader and before returning.

Suggested change

	ralloc_free(b.shader);
	return left == 0 ? agx : 0;
	ralloc_free(b.shader);
	glsl_type_singleton_decref();
	return left == 0 ? agx : 0;

[maralcbr]

Thanks @joshuaswarren, the fix is right. The lowering already widens the accumulator to f32 (agx_nir_lower_cmat.c:249-250) and converts the result back to the result type before the store (:256-257), so only the assertion at :240-246 was in the way. I found no blockers in the code. Builds with assertions now lower this shader instead of aborting. Release builds already compiled this path the same way, and the advertised configurations in hk_GetPhysicalDeviceCooperativeMatrixPropertiesKHR don't change.

Smaller things:

• Validate in the test. lowered_muladd_count calls the pass directly (test-cmat-accumulator.cpp:58), and nir_progress doesn't validate, so a malformed store would still pass the count check. Please call nir_validate_shader(b.shader, "after agx_nir_lower_cmat") after the pass. Ideally also check that an f2f16 feeds the final store in the half case, so the test covers the conversion back too.
• A question rather than a fix: since the driver doesn't advertise an f16 accumulator for f16 inputs (as the body says), do you know whether llama.cpp gates that flash-attention path on the reported cooperative matrix properties? It would tell us whether there's something to report upstream.

Nits, fine as follow-ups:

• Please add glsl_type_singleton_decref() after ralloc_free (test-cmat-accumulator.cpp:74), as test-demote-early-kill.cpp:223-227 does.
• The assertion now also admits an f16 accumulator with bf16, mixed f16/bf16 or 8x8 f32 inputs. That's fine because the conversion is generic, but a short comment saying the accumulator is always widened to f32 would help.

CI: the x86_64 and aarch64 build jobs were still pending when I looked, so there's no completed agx_tests result at this head yet. CI builds with -Db_ndebug=false, so the new test will exercise the assertion there.

Whether to advertise f16 x f16 -> f16 configurations is a maintainer call for a separate PR, not this one. For release (not the source merge): the llama.cpp flash-attention shader still needs a run on the device with this change, per the edge qualification in AGENTS.md.

Reviewed at a8b0011539c2.

— Marcelo's agent

[maralcbr]

Could you add nir_validate_shader(b.shader, "after agx_nir_lower_cmat"); here? nir_progress doesn't validate, so a mis-sized store would still pass the count.

— Marcelo's agent

[maralcbr]

nit: glsl_type_singleton_decref(); after ralloc_free, to balance the ref taken at line 41.

— Marcelo's agent

[maralcbr]

nit: worth a short comment that C is always widened to f32 (line 249), so any float accumulator type is fine here.

— Marcelo's agent