Save size in scalar scratch for bo and bq

[rupengliu-meta]

Description

In bo and bq, we could save size in smem to avoid calculation. This will reduce unnecessary computation.
Tests have passed for both kernels

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

Tests

Ran unit tests and done local e2e testing
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[yaochengji]

Thanks for the contribution! I think the trade-off is between scalar computation and scalar load/store, do you have any performance number after the modification?

[rupengliu-meta]

Thanks for the contribution! I think the trade-off is between scalar computation and scalar load/store, do you have any performance number after the modification?

yes, I will update the perf numbers later

[rupengliu-meta]

seems only having pretty minimal throughput improvement, but the improvement is consistently around 1%-2%. tested through the kernel benchmarking script (not e2e)

[vanbasten23]

do you know why queries is not in the donate_argnames in jax.jit?

[rupeng-liu]

it depends on if it will be used again after the attention. I took a quick look and didn't find where it is used again. So if no reuse of the queries, then we could donate it. @bythew3i thought?

[kyuyeunk]

Isn't this change also applicable for bkv as well? i.e., save bkv sz to a scalar scratch?

[rupeng-liu]

@kyuyeunk yep, good idea. I just checked the bkv sz, the sz is offset + bkv_sz_frm_new, which during wait is False, there is no existing value for this, we need to still do the extra calculation if added in the wait=false. So this might not be applicable for bkv?

[kyuyeunk]

Thanks for the contribution! I think the trade-off is between scalar computation and scalar load/store, do you have any performance number after the modification?

yes, I will update the perf numbers later

Ping on updating perf numbers on the pr description.