Gemma 4 follow-ons / clean-ups

@hunhoffe

LM Head needs to use shared testing infrastructure rather than rolling its own.
Unify Q4NX helpers for LM Head, dequant, decode, attention operators
Make sure all dispatch parameters, such as context_len, are validated before the dispatch for design-valid values.
Low-hanging-fruit cleanups for the decode operator:
- Get rid of the _Ctx struct, make the builder functions simple functions rather than passing that struct around.
- Remove duplicated constants, such as RMS_EPS, into a single source definition. The kernel factory in MLIR-AIE might need to export them, necessitating an MLIR-AIE PR.
iron/operators/flm/aie2p_math_emulation.py -- a lot of these already exist in MLIR-AIE according to @hunhoffe; find and replace usages in IRON to remove duplicated code; if some changes or additions to the helpers in MLIR-AIE are required, that is acceptable and preferable over duplicating this code here

Status check on all five items — 4 are done, narrowing this to the last one.

LM Head shared testing infra — done, flm/lm_head/test.py uses the shared Testing/run_test/vectors harness.
Q4NX unification — done, one shared iron/operators/flm/q4nx.py.
Dispatch parameters validated before dispatch — done. iron/operators/flm/layer/design.py validates context_len >= 0, max_l <= MAX_CONTEXT, and context_len < max_l at the top of the runtime-sequence builder, before any RTP writes or shim DMA starts.
_Ctx struct / RMS_EPS duplication — done, both cleaned up.
Dedup aie2p_math_emulation.py against mlir-aie — still open

Navigation Menu