Fixing autoregressive generation for local sliding window and chunk attention

[gagika]

Description

Fixing autoregressive generation for local sliding window and chunk attention, for both cases using query offset to compute the mask

Tests

Tested attention masks for local and Llama4 autoregressive generation.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[RissyRan]

LGTM! Wondering if we could attach the test you have run?

Also, @aireenmei could you remind which unit test of attention is failing? is that MLA?

[gagika]

LGTM! Wondering if we could attach the test you have run?

Also, @aireenmei could you remind which unit test of attention is failing? is that MLA?

@RissyRan

Llama4 test for autoregressive generation: b/415126351#comment3

This fix is not related to MLA, but wondering what was the issue?

[RissyRan]

LGTM! Wondering if we could attach the test you have run?
Also, @aireenmei could you remind which unit test of attention is failing? is that MLA?

@RissyRan

Llama4 test for autoregressive generation: b/415126351#comment3

This fix is not related to MLA, but wondering what was the issue?

Yeah, a separate issue I think. She mentioned MLA unit test does not pass, we could take this offline.

Testing in https://b.corp.google.com/issues/415126351#comment3, LGTM! Thank you!

[khatwanimohit]

LGTM