New custom op for KV cache update

[vineet-g]

Summary

• Unify KV-cache update + fetch into a single fused custom op
• Both iOS and macOS model definitions now use the new custom op

[carinapeng]

Thank you! A high level question for the motivation of this fused op, do we have perf numbers? Or is this for correctness?

Am I right in the understanding that at the MLIR level the op is decomposed away into unsqueeze -> immutable_slice_update -> aten.slice -> aten.slice -> aten.squeeze?

Can we get the same primitives from slice_update + slice + squeeze primitives? Or is the intention to change dependency as we lower?

[TobyRoseman]

This looks interesting. Would it be possible to add unit tests for this new functionality?

[vineet-g]

Thank you! A high level question for the motivation of this fused op, do we have perf numbers? Or is this for correctness?

This is to enable KV-cache quantization. No perf change is expected; the lowered MLIR has the same primitives as before.

Am I right in the understanding that at the MLIR level the op is decomposed away into unsqueeze -> immutable_slice_update -> aten.slice -> aten.slice -> aten.squeeze?

Yes

Can we get the same primitives from slice_update + slice + squeeze primitives? Or is the intention to change dependency as we lower?

Existing mutable_slice_update takes one layer's K/V in and returns the whole cache out which creates asymmetry while placing quantize/dequantize nodes around the custom op.

The new fused op takes one layer in and returns that same layer's populated cache as output which gives the quantizer a symmetric custom-op boundary.

[carinapeng]

Looks good, but let's make sure we are not introducing regression with the python tests, and +1 to @TobyRoseman 's point on adding unit test for this

[vineet-g]

Would it be possible to add unit tests for this new functionality?

python/tests/test_model_units/test_primitives/test_macos/test_cache.py already covers the unit tests for cache_update which now uses the new custom op