Update silu_and_mul.py to fix the int32 overflow issue (#354)

ModelTC · Mar 14, 2024 · e2b5168 · e2b5168
1 parent 5109f0b
commit e2b5168
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/lightllm/models/llama/triton_kernel/silu_and_mul.py b/lightllm/models/llama/triton_kernel/silu_and_mul.py
@@ -16,6 +16,9 @@ def _silu_and_mul_kernel(
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
 ):
+    stride_input_m = stride_input_m.to(tl.int64)
+    stride_output_m = stride_output_m.to(tl.int64)
+
     tid = tl.program_id(0)
     input_m_offsets = tid * BLOCK_M + tl.arange(0, BLOCK_M)
     output_m_offsets = tid * BLOCK_M + tl.arange(0, BLOCK_M)