diff --git a/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h b/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h index af7a202ef559f..88a22fd422af0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h +++ b/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h @@ -476,13 +476,17 @@ struct AttentionKernel { kNumWarpsPerBlock, ""); - // used for efficient load of bias tile Bij from global to shared memory + // used for efficient load of bias tile Bij from global to shared memory. + // Use kAlignmentA so the unaligned kernel path (kIsAligned=false) uses + // narrower vectorized loads (64-bit instead of 128-bit), matching the + // relaxed alignment requirement for Q/K/V. This allows bias_strideM + // (= total_kv_length) to be any multiple of 4 elements (fp16) rather + // than requiring a multiple of 8. using BiasLoader = TileSmemLoader< scalar_t, cutlass::MatrixShape, MmaCore::kThreads, - // input restriction: kv_len has to be a multiple of this value - 128 / cutlass::sizeof_bits::value>; + kAlignmentA>; // Epilogue to store to shared-memory in a format that we can use later for // the second matmul diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/test_mha.py b/onnxruntime/test/python/transformers/test_onnx_attention/test_mha.py index a488e11e39d20..271d66190c9e7 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/test_mha.py +++ b/onnxruntime/test/python/transformers/test_onnx_attention/test_mha.py @@ -2531,6 +2531,127 @@ def test_flash_softcap_with_mask_prompt_fp16(self): ) +@unittest.skipIf(not has_cuda_device(53), "CUDA device not available, skipping MHA tests.") +@patch.dict(os.environ, {"ORT_DISABLE_FLASH_ATTENTION": "1"}) +class TestONNXAttentionMHACutlassBiasAlignment(unittest.TestCase): + """Test CUTLASS BiasLoader alignment with unaligned total_kv lengths. + + CUTLASS Memory Efficient Attention (MEA) uses vectorized loads for the + attention bias. Before the fix in PR #28369, the BiasLoader hardcoded + 128-bit (8 x fp16) alignment, causing wrong results or crashes when + total_kv_length % 8 != 0. The fix uses kAlignmentA which respects the + unaligned kernel path. + + These tests verify that both aligned (% 8 == 0) and unaligned (% 8 != 0) + sequence lengths produce correct results with additive float masks. + """ + + @parameterized.expand( + [ + ("unaligned_5", 4, 5), + ("unaligned_7", 6, 7), + ("unaligned_9", 8, 9), + ("unaligned_13", 12, 13), + ("unaligned_27", 26, 27), + ("aligned_8", 7, 8), + ("aligned_16", 15, 16), + ("aligned_32", 31, 32), + ] + ) + def test_mha_bias_alignment_decode(self, name, past_kv_len, total_kv_len): + """Decode step (q_seq=1) with additive float mask at various total_kv lengths.""" + config = AttentionConfig( + batch_size=2, + q_sequence_length=1, + kv_sequence_length=1, + past_kv_sequence_length=past_kv_len, + q_num_heads=4, + kv_num_heads=4, + head_size=64, + is_causal=1, + has_attn_mask=True, + attn_mask_dims=4, + attn_mask_type="additive", + ) + parity_check_mha_past( + config=config, + ep="CUDAExecutionProvider", + device="cuda", + torch_type=torch.float16, + ort_type=TensorProto.FLOAT16, + causal=True, + rtol=rtol["fp16"], + atol=atol["fp16"], + ) + + @parameterized.expand( + [ + ("unaligned_5", 5), + ("unaligned_7", 7), + ("unaligned_13", 13), + ("aligned_8", 8), + ("aligned_16", 16), + ] + ) + def test_mha_bias_alignment_prompt(self, name, kv_seq_len): + """Prompt (no past) with additive float mask at various kv_seq lengths.""" + config = AttentionConfig( + batch_size=2, + q_sequence_length=kv_seq_len, + kv_sequence_length=kv_seq_len, + q_num_heads=4, + kv_num_heads=4, + head_size=64, + is_causal=1, + has_attn_mask=True, + attn_mask_dims=4, + attn_mask_type="additive", + ) + parity_check_mha_prompt( + config=config, + ep="CUDAExecutionProvider", + device="cuda", + torch_type=torch.float16, + ort_type=TensorProto.FLOAT16, + causal=True, + rtol=rtol["fp16"], + atol=atol["fp16"], + ) + + @parameterized.expand( + [ + ("unaligned_5", 4, 5), + ("unaligned_9", 8, 9), + ("aligned_16", 15, 16), + ] + ) + def test_gqa_bias_alignment_decode(self, name, past_kv_len, total_kv_len): + """GQA decode (q_heads != kv_heads) with additive mask at various lengths.""" + config = AttentionConfig( + batch_size=2, + q_sequence_length=1, + kv_sequence_length=1, + past_kv_sequence_length=past_kv_len, + q_num_heads=8, + kv_num_heads=2, + head_size=64, + is_causal=1, + has_attn_mask=True, + attn_mask_dims=4, + attn_mask_type="additive", + ) + parity_check_mha_past( + config=config, + ep="CUDAExecutionProvider", + device="cuda", + torch_type=torch.float16, + ort_type=TensorProto.FLOAT16, + causal=True, + rtol=rtol["fp16"], + atol=atol["fp16"], + ) + + # NOTE: GQA fully-masked batch fix (ZeroOutputForFullyMaskedBatches) is validated by # C++ test Attention_NonPadKVSeqLen_AllMasked_FP16_GQA. Python graph-level test omitted # because the fix is a CUDA kernel in the MEA path — a CPU-only test cannot validate it,