|
55bc58a842
|
fix(triton): widen blocked RMSNorm row index (#1482)
## Summary
Widen the blocked RMSNorm forward program ID before row-offset address
arithmetic. The default blocked path (`n_cols <= 256`, `n_rows >=
32768`) otherwise overflows a 32-bit row address at 2^31 elements.
This completes the same safeguard introduced for the non-blocked RMSNorm
kernels in #804.
Fixes #1483.
## Testing Done
- Hardware Type: NVIDIA RTX 2000 Ada Generation Laptop GPU (8 GB)
- [x] `python -m pytest -q test/transformers/test_rms_norm.py -k
block_rms_norm -x` (1 passed; large CUDA regression skipped because it
requires 8.6 GB free VRAM)
- [x] `ruff check src/liger_kernel/ops/rms_norm.py
test/transformers/test_rms_norm.py`
- [x] `ruff format --check src/liger_kernel/ops/rms_norm.py
test/transformers/test_rms_norm.py`
- [ ] `make test` (make is unavailable on this host)
- [ ] `make checkstyle` (make is unavailable on this host; its ruff
commands above passed)
- [ ] `make test-convergence` (not run)
Co-authored-by: Srijan Keshri <[email protected]>
Co-authored-by: Jari Kolehmainen <[email protected]>
|
2026-09-24 23:50:30 +00:00 |
|