Add support for unpadded shapes in Matmul1D w/ gather_in0

[avoraTT]

Ticket

• Support unpadded shapes in matmul1d w/ gather_in0 #16626

Problem description

In the current use case of Matmul1D with gather_in0 in the Llama models, the activations and weights need to be padded. This results in significant overhead.

What's changed

To support unpadded shapes, the implementation leverages the implicit padding information that can be stored in the sharding spec. For example, a tensor might be of shape [1, 1, 32, 2048] and we want to shard it over 24 cores. Since K = 2048 doesn't divide with 24 cores, we can pad the shape to K_pad = 2304, while keeping the actual tensor shape the same. Then, the shapes in the sharding spec can be derived using K_pad, instead of K.

Since the matmul block sizes are derived using the sharding spec, the kernels implicitly process the padding shape, where some pages at the end of the inner dim (for activations) contain garbage data. Note, if the activations are being implicitly padded, to make sure the the output is correct when implicitly padding the inner dim, the inner dim for the weights must be explicitly padded with 0's. This is to ensure that when accumulating partials, the garbage/padded part of the activations get multiplied with 0's before summation.

Example for matmul on 24 cores
Original shapes

• activations: [1, 1, 32, 2048]
• weights: [1, 1, 2048, 3584]

Actual tensors passed into the op

• activations: [1, 1, 32, 2048], with shard shape using padding: [32, 2304 / 24]
• weights: [1, 1, 2304, 3584], with shard shape using padding: [2304, 3840 / 24]
• output: (shape internally derived by matmul) [1, 1, 32, 3584], with shard shape using padding: [32, 3840 / 24]

Other changes include moving validation code in matmul to support different K's for the activation vs weights when gather_in0 = True.

Checklist

• Post commit CI passes
• Blackhole Post commit (if applicable)
• Model regression CI testing passes (if applicable)
• Device performance regression CI testing passes (if applicable)
• (For models and ops writers) Full new models tests passes
• New/Existing tests provide coverage for changes

[yugaoTT]

why is it changed from 1 to 0 (from K to M)?

[yugaoTT]

I assume this is a redundant check and there's other places that check for a_shape[-1] == b_shape[-2] ?

[yugaoTT]

commented out?

[yugaoTT]

is this flag essentially program_config.gather_in0? might be confusing for people without context