fix(grpo): use 64-bit row offsets in the GRPO loss kernels

[Lfan-ke]

Problem

grpo_kernel and grpo_bwd_kernel address the logits with pid * V, where pid comes from tl.program_id and is int32. The product wraps once BL * V exceeds 2^31, so the kernels read — and in the backward pass write — outside the tensors. The launch is grid = (BL,) with BL, V = new_logits.shape, so nothing else bounds it.

The threshold is reachable in ordinary use. With a 128k vocabulary it is crossed at 16744 tokens, and grpo_loss_triton_optimized already defaults max_completion_length to 8192.

Fix

Promote pid to int64 in both kernels, so every derived offset is 64-bit. This is the same fix already applied to permute_bwd_kernel in #172.

Also corrected the import in benchmarks/grpo.py, which named the package dlblas.kernel instead of dlblas.kernels and therefore could never run.

Verification

The overflow itself needs ~8.6 TB of logits to reproduce with real tensors, so the address arithmetic was reproduced on its own, with the same expression the kernel uses (RTX 3050, triton 3.7.1):

V=128256  BL=20000  int32 limit reached at pid=16743
rows with a wrong int32 offset : 3256 / 20000
  first wrong pid=16744: got -2147448832, expected 2147518464  (delta -4294967296)
int64 offsets all correct      : True

A negative offset is an out-of-bounds access; in grpo_bwd_kernel it is an out-of-bounds store.

No behaviour change below the threshold — tests/kernels/test_grpo_loss_logits.py gives a bit-identical loss before and after:

PyTorch Loss: 1.6318347454071045
Triton Loss:  1.6318347454071045
Forward Pass (Loss) Correct: True
Backward Pass (Gradient) Correct: True

[CLAassistant]

All committers have signed the CLA.