diff --git a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py
index b71fc97db..80bb66e7e 100644
--- a/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py
+++ b/python/sglang/srt/batch_invariant_ops/batch_invariant_ops.py
@@ -309,9 +309,9 @@ def matmul_persistent(
 def _log_softmax_kernel(
     input_ptr,
     output_ptr,
-    input_row_stride,
-    output_row_stride,
-    n_cols,
+    input_row_stride: tl.constexpr,
+    output_row_stride: tl.constexpr,
+    n_cols: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
 ):
     """
@@ -814,9 +814,9 @@ def _rms_norm_kernel(
     input_ptr,
     weight_ptr,
     output_ptr,
-    input_row_stride,
-    output_row_stride,
-    n_cols,
+    input_row_stride: tl.constexpr,
+    output_row_stride: tl.constexpr,
+    n_cols: tl.constexpr,
     eps,
     BLOCK_SIZE: tl.constexpr,
 ):