From fb609669ca1288f15fcbf9dcd563f61ea9521811 Mon Sep 17 00:00:00 2001
From: Lianmin Zheng <lianminzheng@gmail.com>
Date: Sun, 1 Feb 2026 07:39:37 -0800
Subject: [PATCH] [Auto Sync] Update elementwise.py (20260131) (#18033)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 python/sglang/srt/layers/elementwise.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/sglang/srt/layers/elementwise.py b/python/sglang/srt/layers/elementwise.py
index dd07fc485..d8f2f7e48 100644
--- a/python/sglang/srt/layers/elementwise.py
+++ b/python/sglang/srt/layers/elementwise.py
@@ -230,7 +230,7 @@ def fused_rmsnorm_kernel(
     hidden_dim: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
-    pid = tl.program_id(axis=0)
+    pid = tl.program_id(axis=0).to(tl.int64)
     input_start = pid * hidden_dim
 
     offsets = tl.arange(0, BLOCK_SIZE)