SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

fix(norm): widen address arithmetic to int64 for large contiguous inputs > 2**31 elements (#3392)

<!-- .github/pull_request_template.md -->

## ๐Ÿ“Œ Description

<!-- What does this PR do? Briefly describe the changes and why theyโ€™re
needed. -->

Fix cudaErrorIllegalAddress in the cute-DSL norm kernels (`rmsnorm`,
`gemma_rmsnorm`, `rmsnorm_quant`, `fused_add_rmsnorm`,
`gemma_fused_add_rmsnorm`, `fused_add_rmsnorm_quant`, `layernorm`, and
`qk_rmsnorm`) when the input tensor's flat element count exceeds
`2**31`.

**Root Cause**
The compact (contiguous) compile path bakes the row stride in as a
Python constexpr integer, and M / B / N were declared as `Int32`
everywhere. The row coord materialized into the offset comes from
`cute.arch.block_idx()`, which is intrinsically `int32`, so the compiler
emits `row * H` in `int32` and overflows whenever `M * H > INT32_MAX`.
For shape (175000, 12288) fp16 (174999 * 12288 โ‰ˆ 2.15e9) this manifests
as `cudaErrorIllegalAddress` on the first `cudaStreamSynchronize`. The
strided compile path was already safe because its row stride was already
a dynamic `sym_int64`, but the contiguous path was not.

A secondary bug existed in `flashinfer/norm/utils.py::get_ptr_as_int64`,
which truncated its offset parameter to Int32 before adding it to the
base pointer โ€” so even a correctly-computed `int64` element offset was
being downcast in the rmsnorm-quant kernels' manual FP8 store path.

## ๐Ÿ” Related Issues

<!-- Link any related issues here -->

#3391

## ๐Ÿš€ Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### โœ… Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## ๐Ÿงช Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->


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

* **Bug Fixes**
* Resolved integer-overflow issues in normalization (RMSNorm, LayerNorm,
fused Add+RMSNorm and FP8-quantized variants) for very large tensors,
preventing illegal memory accesses and ensuring correct addressing for
huge row counts.
* Forced a safe compilation/execution path for contiguous inputs when
dimensions exceed 32-bit limits.

* **Tests**
* Added regression tests that detect contiguous overflow scenarios
across normalization variants, with VRAM-aware gating and spot-check
validation.

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3392?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
B
Brian K. Ryu committed
96ab2cedc0f35def8929d12f831d299493ad261c
Parent: d53f106
Committed by GitHub <noreply@github.com> on 5/26/2026, 8:49:17 PM