JAXBench: Fix ragged paged attention benchmark inputs

[charleshong3]

Summary

Fixes the 7p_Ragged_Paged_Attention benchmark, whose previous inputs were nominally ragged but actually used:

• 64 uniformly active sequences
• 64 query tokens per sequence
• identical 4096-token KV lengths
• contiguous identity page mappings

This allowed implementations to treat the workload as fixed-shape dense attention without honoring the ragged metadata.

Changes

• Introduce one documented Llama-3.1-70B continuous-batching scenario:
  • 48 one-token decode sequences
  • 8 chunked-prefill sequences
  • 4096 total query tokens
  • varied KV lengths from 257 to 4095
  • shuffled, noncontiguous physical page mappings
  • 56 active and 8 inactive sequence slots
• Update the XLA reference to honor packed cu_q_lens, variable KV lengths, and inactive sequences while remaining JIT-compatible.
• Zero inactive Pallas output rows instead of exposing uninitialized VMEM contents.
• Use the current pl.ANY API.
• Calculate FLOPs from effective query and KV lengths.
• Add a JAXBench-local JAX 0.9.2 dependency pin without changing dependencies for other Accelerator Agents components.
• Explicitly document that this benchmark measures one representative serving scenario, not the full RPA input distribution.

Validation

Tested on TPU v6e-1 with JAX 0.9.2:

• Full baseline/optimized output match at JAXBench harness tolerance
• Maximum absolute difference: 0.00390625
• Global relative difference: 0.00667
• XLA baseline median: 16.12 ms
• Pallas optimized median: 3.46 ms
• Speedup: 4.66×
• Shuffled page mapping confirmed
• Python compilation and git diff --check pass

[shangkunwang01]

Could you remove this.

[shangkunwang01]

Could you remove this as well.

[shangkunwang01]

Do you have correctness check for this optimized code
Not equal to tolerance rtol=0.1, atol=0.1 Outputs from baseline.py and optimized.py do not match. Mismatched elements: 29672584 / 33554432 (88.4%) First 5 mismatches are at indices: [0, 0, 0]: 0.05458643287420273 (ACTUAL), 0.414387583732605 (DESIRED) [0, 0, 1]: 0.00877971388399601 (ACTUAL), 0.5063981413841248 (DESIRED) [0, 0, 2]: 0.04373247176408768 (ACTUAL), -0.3895590305328369 (DESIRED) [0, 0, 3]: 0.14047656953334808 (ACTUAL), 1.2859950065612793 (DESIRED) [0, 0, 4]: -0.15021640062332153 (ACTUAL), -0.5560941696166992 (DESIRED) Max absolute difference among violations: 5.3871675 Max relative difference among violations: 199319.38 ACTUAL: array([[[ 0.054586, 0.00878 , 0.043732, ..., -0.077147, 0.013508, -0.088013], [-0.08656 , 0.024763, 0.032719, ..., 0.052373, 0.129617,... DESIRED: array([[[ 0.414388, 0.506398, -0.389559, ..., -0.658711, -0.285335, 0.094416], [ 0.085367, -0.116005, 0.952759, ..., 0.557319, 0.065653,...
I found it could not pass the tol=0.1 test.

[charleshong3]

Hmm did you use both my new baseline and optimized code or are you comparing against the old baseline? If this is a comparison against the old baseline, it should not be expected to match as the kernel was updated.

[shangkunwang01]

Could you remove this line.