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 --> [](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