Add custom masking to fused attention

[pearlli98]

I would like to add a custom mask of shape (batch_size, head_dim, seq_length, seq_length) inside the kernel. I am not sure how to find the appropriate range of pointers to the mask inside the for loop. Can anyone offer some insights on this? Thank you!

[pearlli98]

I witnessed that the kernel would output nan after I added a custom mask. The code is copied below for reference. Mask passed in is of shape (bsz, n_heads, seq_len, seq_len), which correspond to the arg stride_dz,stride_dh,stride_d3,stride_d4, in the code below.

would someone be able to point out the issue? Thank you!

start_m = tl.program_id(0)
off_hz = tl.program_id(1)
# initialize offsets
offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
offs_n = tl.arange(0, BLOCK_N)
offs_d = tl.arange(0, BLOCK_DMODEL)
...
off_mask = (
    off_hz * stride_dh + offs_m[:, None] * stride_d3 + offs_n[None, :] * stride_d4
)
# Initialize pointers to Q, K, V
...
mask_ptrs = D + off_mask
# initialize pointer to m and l
t_ptrs = TMP + off_hz * N_CTX + offs_m
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
# load q: it will stay in SRAM throughout
q = tl.load(q_ptrs)
# loop over k, v and update accumulator
for start_n in range(0, N_CTX, BLOCK_N):
    start_n = tl.multiple_of(start_n, BLOCK_N)
    # -- compute qk ----
    k = tl.load(k_ptrs + start_n * stride_kn)
    mask = tl.load(mask_ptrs + start_n)
    # drop = tl.load(drop_ptrs + start_n)
    qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
    qk = tl.dot(q, k, trans_b=True)
    qk *= sm_scale
    qk += tl.where(mask, float("-inf"), 0)