[Splash Attention] Remove Unnecessary head_dim_v Constraint and Update Scratch Array Shapes #27427

[SwarnimShekhar]

This PR addresses issue #27427 by removing the unnecessary constraint that enforced head_dim_v to be a multiple of NUM_LANES. The following changes are made:

Constraint Removal:
The check:

head_dim_v_repeats, rem = divmod(head_dim_v, NUM_LANES)
if rem != 0:
    raise NotImplementedError(
        f"{head_dim_v=} should be a multiple of {NUM_LANES}"
    )

has been removed. This constraint prevented small models (e.g. sLLMs) from using Splash Attention when head_dim_v is not a multiple of NUM_LANES.

Scratch Array Shape Update:
The scratch arrays m_scratch and l_scratch are now allocated with shape [bq, 1] instead of [bq, NUM_LANES], reducing redundant memory allocation. Broadcasting is now relied upon to handle the expansion as needed in downstream operations.

Alpha Repetition Removal:
All occurrences where pltpu.repeat(alpha, head_dim_v_repeats, axis=1) was used have been replaced with a direct assignment (alpha_o = alpha), which leverages broadcasting for correct behavior.

Testing:

All existing tests pass on both parallel and serial test runs.

[ds-hwang]

Lovely. Thank you for quick action!

[SwarnimShekhar]

Glad to contribute!