[EmmetAVS]

Problem

Fixes #816.
When running a quantized (2-bit or 3-bit) model on Metal, the first token outputted is garbage, while running the same model on CPU or a quantized 4-bit model on Metal work properly.

When the final prefill token is being processed on 2/3-bit quantization (with prefill_consistent set), it calls cactus_metal_encode_quant_matmul_m, which rejects 2/3-bit quantization so the matmul falls back to CPU (as cactus_metal_encode_quant_matmul_m only accepts 4-bit quantization). Because the final RMS_NORM in the fused kernel (rule 3) running before the matmul was labeled -2 (covered by cluster), metal_live is not set and therefore Metal is not synced before the CPU matmul is run. As a result, the CPU reads from the output buffer of the RMS_NORM node before it had been written to. Chunked prefill is not affected by this issue because it is multi-row, which doesn't trigger rule 3 and the existing norms run as normal Metal nodes (-1) and properly trigger metal_live. Additionally, decode is not affected because prefill_consistent is set to false, so it uses cactus_metal_encode_quant_matmul, which is able to run on Metal for 2/3-bit quantized models.

Solution

Modified one line in metal_plan.cpp, specifying that when rule 3's fused kernel includes the next norm (c.u0 == 1), that node (c.a5) is marked as action = -3 (read outside cluster) and therefore triggers metal_live to be set. Followed the standard set by other existing rules (12, 16) where after add_cluster succeeds, covered (-2) nodes read outside their cluster are marked -3.

Verification

• Created parity_cpu_reader_after_fused_norm() unit test in cactus-graph/tests/test_metal_parity.cpp.
• The unit test builds the pattern that triggers rule 3 (rms_norm -> add_clipped -> view -> rms_norm -> reshape -> view -> matmul). It then forces the final matmul to ComputeBackend::CPU and runs that graph on both CPU and Metal and confirms the outputs are within tolerance.

$ cactus test --suite metal_parity    # without fix
✗ FAIL │ cpu_reader_after_fused_norm

$ cactus test --suite metal_parity    # with fix
✓ PASS │ cpu_reader_after_fused_norm
✓ All 48 tests passed!

• On Gemma 4 E2B CQ2 and CQ3, Metal first tokens now match CPU across 4 prompts

Notes

• In parity_cpu_reader_after_fused_norm, c.check() was not used because it does not properly handle NaNs (as comparisons involving NaNs always output false), causing the test to always pass even without the fix (I am happy to make a PR fixing this too)
• The issue notes that "cq3 and cq2 prefill on Metal is about 550 tok/s, basically the same as CPU, while cq4 on Metal does about 1,030", this is because cq3 and cq2 have prefill batched matmul run on CPU because of cactus_metal_encode_quant_matmul_m only accepting 4-bit quantization (not addressed here)