Fix error illegal memory access when max_total_token_num is too large (#998)

Author: sufubao

lightllm/models/bloom/triton_kernel/context_flashattention_nopad.py

@@ -76,7 +76,7 @@ def _fwd_kernel(
             Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * (start_n + offs_n),
             mask=(start_n + offs_n) < block_end_loc,
             other=0,
-        )
+        ).to(tl.int64)
         off_k = kv_loc[None, :] * stride_kbs + cur_head * stride_kh + offs_d[:, None] * stride_kd
         k = tl.load(K + off_k, mask=(start_n + offs_n[None, :]) < block_end_loc, other=0.0)
 

lightllm/models/deepseek2/triton_kernel/context_flashattention_nopad.py

@@ -91,7 +91,7 @@ def _fwd_kernel(
             Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * (start_n + offs_n),
             mask=(start_n + offs_n) < block_end_loc,
             other=0,
-        )
+        ).to(tl.int64)
         off_kv = kv_loc[None, :] * stride_kv_bs + cur_kv_head * stride_kv_h + offs_d[:, None] * stride_kv_d
         off_kv_rope = (
             kv_loc[None, :] * stride_kv_rope_bs

lightllm/models/deepseek2/triton_kernel/context_flashattention_nopad_fp8.py

@@ -96,7 +96,7 @@ def _fwd_kernel_fp8(
             Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * (start_n + offs_n),
             mask=(start_n + offs_n) < block_end_loc,
             other=0,
-        )
+        ).to(tl.int64)
         off_kv = kv_loc[None, :] * stride_kv_bs + cur_kv_head * stride_kv_h + offs_d[:, None] * stride_kv_d
         off_kv_rope = (
             kv_loc[None, :] * stride_kv_rope_bs

lightllm/models/deepseek2/triton_kernel/destindex_copy_kv.py

@@ -34,7 +34,7 @@ def _fwd_kernel_destindex_copy_kv(
     offs_d_nope = tl.arange(0, BLOCK_DMODEL_NOPE)
     offs_d_rope = tl.arange(0, BLOCK_DMODEL_ROPE)
 
-    dest_index = tl.load(Dest_loc + cur_index)
+    dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
 
     kv_nope_ptrs = KV_nope + cur_index * stride_kv_nope_bs + stride_kv_nope_d * offs_d_nope[None, :]
     kv_rope_ptrs = KV_rope + cur_index * stride_kv_rope_bs + stride_kv_rope_d * offs_d_rope[None, :]

lightllm/models/deepseek2/triton_kernel/destindex_copy_kv_fp8.py

@@ -36,7 +36,7 @@ def _fwd_kernel_destindex_copy_kv_fp8(
     offs_d_nope = tl.arange(0, BLOCK_DMODEL_NOPE)
     offs_d_rope = tl.arange(0, BLOCK_DMODEL_ROPE)
 
-    dest_index = tl.load(Dest_loc + cur_index)
+    dest_index = tl.load(Dest_loc + cur_index).to(tl.int64)
 
     kv_nope_ptrs = KV_nope + cur_index * stride_kv_nope_bs + stride_kv_nope_d * offs_d_nope[None, :]
     kv_rope_ptrs = KV_rope + cur_index * stride_kv_rope_bs + stride_kv_rope_d * offs_d_rope[None, :]

lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1.py

@@ -105,7 +105,7 @@ def _fwd_kernel_flash_decode_stage1_padding(
                     req_to_tokens_ptr + offs_n_new,
                     mask=seq_n_mask,
                     other=0,
-                )
+                ).to(tl.int64)
                 off_kv = kv_loc[None, :] * stride_kv_bs + offs_d[:, None]
                 kv = tl.load(KV_nope + off_kv, mask=seq_n_mask[None, :], other=0.0)
                 att_value = tl.dot(q, kv)

lightllm/models/deepseek2/triton_kernel/gqa_flash_decoding_stage1_fp8.py

@@ -108,7 +108,7 @@ def _fwd_kernel_flash_decode_stage1_padding_fp8(
                     req_to_tokens_ptr + offs_n_new,
                     mask=seq_n_mask,
                     other=0,
-                )
+                ).to(tl.int64)
                 off_kv = kv_loc[None, :] * stride_kv_bs + offs_d[:, None]
                 kv = tl.load(KV_nope + off_kv, mask=seq_n_mask[None, :], other=0.0)
                 off_rope_kv = kv_loc[None, :] * stride_kv_rope_bs + offs_rope_d[:, None]

lightllm/models/deepseek2/triton_kernel/sample_kv.py

@@ -44,7 +44,7 @@ def _sample_kv_kernel(
         Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_m,
         mask=offs_m < block_end_loc,
         other=0,
-    )
+    ).to(tl.int64)
     off_kv_nope = kv_loc[:, None] * stride_input_dim + offs_nope_d[None, :]
     off_kv_rope = kv_loc[:, None] * stride_input_dim + (offs_rope_d + BLOCK_DMODEL)[None, :]
     kv_nope = tl.load(KV_input + off_kv_nope, mask=offs_m[:, None] < block_end_loc, other=0.0)

lightllm/models/llama/triton_kernel/context_flashattention_nopad.py

@@ -84,7 +84,7 @@ def _fwd_kernel(
             Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + stride_req_to_tokens_s * (start_n + offs_n),
             mask=(start_n + offs_n) < block_end_loc,
             other=0,
-        )
+        ).to(tl.int64)
         off_k = kv_loc[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None] * stride_kd
         k = tl.load(K + off_k, mask=(start_n + offs_n[None, :]) < block_end_loc, other=0.0)
         qk = tl.dot(q, k)

lightllm/models/llama/triton_kernel/gqa_decode_flashattention_nopad.py

@@ -70,7 +70,7 @@ def _fwd_kernel(
             Req_to_tokens + cur_batch_req_idx * stride_req_to_tokens_b + start_n + offs_n,
             mask=(start_n + offs_n) < cur_batch_seq_len,
             other=0,
-        )
+        ).to(tl.int64)
         k = tl.load(
             k_ptrs + kv_loc[None, :] * stride_kbs, mask=(start_n + offs_n[None, :]) < cur_batch_seq_len, other=0.0
         )