test: pin the int32 page-offset overflow in _page_split_kernel

[edenfunf]

Motivation

_page_split_kernel takes page_idx and sub from tl.program_id(0) (int32) and builds two byte offsets from them:

src_base = src_ptr + page_idx * src_stride0
dst_base = dst_ptr + (page_idx * RATIO + sub) * dst_stride0

Both products are int32. Past 2**31 they wrap, and Triton sign-extends the result onto the 64-bit pointer, so the access lands about 2 GiB below the buffer: a silent write into a neighbouring allocation, or an illegal access if that page is unmapped.

#34025 reported this and #34027 carries the fix. This adds the regression coverage to go with it, so the offset cannot quietly return to int32 in a later refactor.

Modifications

Adds test/registered/kernels/ops/attention/test_page_split_int32_offset.py on base-b-kernel-unit / 1-gpu-large.

The two offsets wrap at different pool sizes, so each gets its own test:

• destination: dst page 57,359, reached once the pool holds 14,340 source pages (3,671,040 tokens/rank)
• source: page_idx 14,340, at the shipped 149,760 B page stride

Both strides are kernel parameters, so each test packs the axis it is not exercising. Only one buffer has to span 231 at a time, which holds peak at 4.13 GiB instead of ~8. The source test sets src_stride0 to 227 so the wrap arrives at page 16 and its destination stays a few MB.

Each large buffer sits above a canary region in the same allocation. A regression then writes into memory the test owns and trips an assertion, rather than raising an illegal access that would take the CUDA context down for everything else in the process.

Masked page skipping, the alignment tail and the persistent-buffer contract are already covered by TestTouchedPageSplit in test_flash_mla_backends.py, through the production entry point. This file only covers the arithmetic past 2**31.

Includes the one-line fix from #34027 so the test passes standalone. I will drop it and rebase if that lands first.

Accuracy Tests

RTX 5070 (SM120), torch 2.9.1+cu128 / triton 3.5.1. Peak 4.13 GiB, nothing retained between tests.

• Reverting the fix fails both tests, 6 assertions, no illegal access. The canary catches the destination wrap; the source wrap reads canary bytes into the destination, where the signature check catches it.
• Mutation check. Setting data_src_off = 0 fails 53 assertions. Source pages carry a per-(page, sub) signature so that class of regression cannot hide behind a uniform fill.
• Boundary arithmetic, checked separately. Replicating the offset expressions and widening after the multiply puts the int32 wrap at -2,147,446,336 (pid 57,359, destination) and -2,147,408,896 (pid 57,360, source). int64 matches exact arithmetic at every pid.
• Well past the boundary. A 100,000-page run at the production stride, covering 42,641 wrapping destination programs and 10,660 wrapping source pages, is byte-exact throughout.

Speed Tests and Profiling

At 8,192 pages, where neither variant wraps: int32 4.818 ms, int64 4.797 ms (-0.43%, noise). One widened scalar op against ~37 KB of copy per program.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.