SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

fix(gdn_decode): widen pool indices to Int64 to prevent int32 element-offset overflow (#3230)

## ๐Ÿ“Œ Description

Fix `CUDA error: an illegal memory access was encountered` in
`flashinfer.gdn_decode.gated_delta_rule_decode_pretranspose` when the
pool+indices API is used with sufficiently large pool indices.

**Root cause.** The CuTe-DSL kernels compute the per-slot element offset
(`pool_idx * stride[0]`, or `(cache_idx * HV + i_hv) * stride[0]` for
bf16) using **Int32** arithmetic. Once the product exceeds `INT32_MAX`,
it wraps to a negative offset and the load/store hits an unmapped global
address.

Affects both backends the API can dispatch to (HV=32, V=K=128):

| backend | kernel | overflow threshold |
|---|---|---|
| fp32 pretranspose | `gdn_decode_kernel_{small,big}_batch_pretranspose`
| `pool_idx >= 3972` (vLLM padded slot stride 540 672) |
| bf16 fast path | `gdn_decode_bf16state_mtp_kernel` | `cache_idx >=
4096` (contiguous, `stride[0] = HV*V*K = 524 288`) |

Discovered while integrating the kernel into vLLM's GDN decode path for
**Qwen3.5-class models**.

**Fix.** Widen the pool indices to Int64 immediately after they are
read; downstream offsets in `cute.local_tile(...)` / `h0_source[(...)]`
then promote to Int64 and cannot wrap:

```python
# fp32 pretranspose (small + big batch)
pool_idx = cutlass.Int64(h0_indices[i_n])
out_pool_idx = cutlass.Int64(h0_out_indices[i_n])

# bf16 MTP โ€” propagates Int64 through flat_state_idx,
# flat_write_idx, and the intermediate-states cache's flat_idx.
cache_idx = cutlass.Int64(h0_indices[i_n])
write_cache_idx = cutlass.Int64(h0_out_indices[i_n])
```

## ๐Ÿ” Related Issues

Same class of bug as
[#3005](https://github.com/flashinfer-ai/flashinfer/issues/3005) /
[#3007](https://github.com/flashinfer-ai/flashinfer/pull/3007) (rmsnorm
stride overflow), in a different family of CuTe-DSL kernels.

## ๐Ÿš€ Pull Request Checklist

### โœ… Pre-commit Checks

- [x] `pre-commit` installed and hooks installed.
- [x] `pre-commit run --files <changed files>` โ€” all hooks pass.

## ๐Ÿงช Tests

- [x] Tests added.
- [x] All tests pass.

Added `tests/gdn/test_decode_pretranspose_noncontiguous_pool.py`:

- `test_decode_pretranspose_pool_int64_offset[3972, 8191]` โ€” fp32
vLLM-padded pool (~8.6 / 17.7 GB).
- `test_decode_pretranspose_pool_int64_offset_bf16[4096, 4196]` โ€” bf16
contiguous pool (~4.3 GB).

Both compare the pool path against a gather + direct-state reference
(numerical correctness, not just non-crashing) and assert the in-place
state update matches. VRAM-based skip when free memory is insufficient.

**Verified on NVIDIA B200 (SM100):** all 4 new tests crash without the
fix and pass with it; existing `pretranspose` and `bf16_state` tests in
`tests/gdn/test_decode_delta_rule.py` continue to pass.



<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Resolved integer overflow in GPU decode kernels by switching pool- and
state-index arithmetic to 64-bit, preventing wraparound and
out-of-bounds addressing for large pools and batches.
* Ensured consistent 64-bit handling across all decode paths and
negative-index clamping.

* **Tests**
* Added GPU regression tests covering large-pool overflow scenarios for
FP32 and BF16 fast paths, with device-capacity guards to avoid OOM on
low-VRAM systems.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
V
Vadim Gimpelson committed
bc39fbc058973a7c71818daf6c261e21ed91b0bd
Parent: 4381afc
Committed by GitHub <noreply@github.com> on 5/11/2026, 11:48:43 PM