[Feature]: Fix the mandatory boundary_check when loading bias tensor

[xinyazhang]

Suggestion Description

The boundary_check option should only be applied when seqlen_q/k cannot be divided by BLOCK_M/N

aotriton/tritonsrc/bwd_split_kernel.py

Lines 218 to 229 in d6c6dcc

	# FIXME: do boundary_check correctly
	"""
	if q_padded and k_padded: # CAVEAT: using "or" disables the partial boundary_check branches
	bias = tl.load(B_block_ptr, boundary_check=(0,1), padding_option="zero")
	elif q_padded:
	bias = tl.load(B_block_ptr, boundary_check=(0,), padding_option="zero")
	elif k_padded:
	bias = tl.load(B_block_ptr, boundary_check=(1,), padding_option="zero")
	else:
	bias = tl.load(B_block_ptr)
	"""
	bias = tl.load(B_block_ptr, boundary_check=(0,1), padding_option="zero")

aotriton/tritonsrc/bwd_split_kernel.py

Lines 458 to 470 in d6c6dcc

	'''
	if q_padded and k_padded: # CAVEAT: using "or" disables the partial boundary_check branches
	bias = tl.load(B_block_ptr, boundary_check=(0,1), padding_option="zero")
	elif q_padded:
	bias = tl.load(B_block_ptr, boundary_check=(0,), padding_option="zero")
	elif k_padded:
	bias = tl.load(B_block_ptr, boundary_check=(1,), padding_option="zero")
	else:
	bias = tl.load(B_block_ptr)
	'''
	# FIXME: Must use boundary_check uncondtionally.
	# The optimized tl.load above causes nan for some reason
	bias = tl.load(B_block_ptr, boundary_check=(0,1), padding_option="zero")

However according to our tests, it is mandatory for all cases. Hence there are two questions to be answered here

1. Why is this mandatory?
2. Can we remove it?

Operating System

No response

GPU

No response

ROCm Component

No response