bugfix: avoid potential illegal memory access (#267)

yzh119 · web-flow · commit 79a21253403d · 2024-05-27T23:40:16.000-07:00
Followup of #266, add guard to mask array access.
diff --git a/include/flashinfer/attention/prefill.cuh b/include/flashinfer/attention/prefill.cuh
@@ -573,11 +573,11 @@ __device__ __forceinline__ void mask_s(const uint32_t qo_idx_base, const uint32_
                  ? (kv_idx > kv_len + q_idx - qo_len || (partition_kv && kv_idx >= chunk_end))
                  : kv_idx >= chunk_end);
         s_frag[fx][fz][reg_id] =
-            out_of_boundary
-                ? DTypeQKAccum(-5e4)
-                : s_frag[fx][fz][reg_id] + DTypeQKAccum(mask_mode == MaskMode::kCustom
-                                                            ? custom_mask[q_idx * kv_len + kv_idx]
-                                                            : 0.f);
+            out_of_boundary ? DTypeQKAccum(-5e4)
+                            : s_frag[fx][fz][reg_id] +
+                                  DTypeQKAccum((mask_mode == MaskMode::kCustom && q_idx < qo_len)
+                                                   ? custom_mask[q_idx * kv_len + kv_idx]
+                                                   : 0.f);
       }
     }
   }

Original file line number	Diff line number	Diff line change
`@@ -573,11 +573,11 @@ __device__ __forceinline__ void mask_s(const uint32_t qo_idx_base, const uint32_`
`573`	`573`	`? (kv_idx > kv_len + q_idx - qo_len \|\| (partition_kv && kv_idx >= chunk_end))`
`574`	`574`	`: kv_idx >= chunk_end);`
`575`	`575`	`s_frag[fx][fz][reg_id] =`
`576`		`- out_of_boundary`
`577`		`- ? DTypeQKAccum(-5e4)`
`578`		`- : s_frag[fx][fz][reg_id] + DTypeQKAccum(mask_mode == MaskMode::kCustom`
`579`		`- ? custom_mask[q_idx * kv_len + kv_idx]`
`580`		`- : 0.f);`
	`576`	`+ out_of_boundary ? DTypeQKAccum(-5e4)`
	`577`	`+ : s_frag[fx][fz][reg_id] +`
	`578`	`+ DTypeQKAccum((mask_mode == MaskMode::kCustom && q_idx < qo_len)`
	`579`	`+ ? custom_mask[q_idx * kv_len + kv_idx]`
	`580`	`+ : 0.f);`
`581`	`581`	`}`
`582`	`582`	`}`
`583`	`583`	`}`