[Tutorial] Fix 06-fused-attention.py of FP8 provider

[whitneywhtsang]

When the provider is fp8, v is permuted like below, and the new stride is (H*N_CTX*HEAD_DIM, N_CTX*HEAD_DIM, 1, N_CTX).

        if mode == "fwd" and "fp8" in provider:
            v = v.permute(0, 1, 3, 2).contiguous()
            v = v.permute(0, 1, 3, 2)

This PR fixes the FP8 dtype handling in the fused-attention kernel by separating k and v offset calculations and updating related configuration details. Key changes include:

• Renaming and separating offset variables for k and v computations.
• Adjusting offset calculation for FP8 dtype and updating the tensor descriptor creation.
• Expanding configuration options for BLOCK_N and refining device-specific configuration conditions.

[Copilot]

Pull Request Overview

This PR fixes the FP8 dtype handling in the fused-attention kernel by separating key and value offset calculations and updating related configuration details. Key changes include:

• Renaming and separating offset variables for key and value computations.
• Adjusting offset calculation for FP8 dtype and updating the tensor descriptor creation.
• Expanding configuration options for BLOCK_N and refining device-specific configuration conditions.

Comments suppressed due to low confidence (4)

python/tutorials/06-fused-attention.py:55

• [nitpick] The variable 'offsetk_y' now clearly denotes the key tensor offset. Consider updating any adjacent comments or documentation to clarify the separation between key and value offsets to enhance readability.

offsetk_y = offset_y + lo

python/tutorials/06-fused-attention.py:56

• [nitpick] It would be beneficial to add a comment explaining why offsetv_y is computed as 'offset_y * HEAD_DIM + lo' for the FP8 dtype, especially in light of the new stride requirements.

if dtype == tl.float8e5:

python/tutorials/06-fused-attention.py:171

• [nitpick] Consider adding an inline comment to explain the purpose and expected behavior of FP8_OUTPUT in controlling the descriptor configuration, to aid future maintenance.

if FP8_OUTPUT:

python/tutorials/06-fused-attention.py:130

• [nitpick] Consider including a brief comment explaining the significance of checking for device capability 9 and the BLOCK_M * BLOCK_N threshold to improve clarity for future maintainers.

and torch.cuda.get_device_capability()[0] == 9 and BLOCK_M * BLOCK_N < 128 * 128

[peterbell10]

Thanks, can you add an fp8 case to test_op?

[whitneywhtsang]

Thanks, can you add an fp8 case to test_op?

It appears that test_op is written for bwd, and from looking at bench_flash_attention, there are no differences between provider triton-fp16 and triton-fp8 for bwd.

[peterbell10]

It tests both forward and backward:

triton/python/tutorials/06-fused-attention.py

Lines 613 to 620 in 5450875

	# triton implementation
	tri_out = attention(q, k, v, causal, sm_scale, warp_specialize).half()
	tri_out.backward(dout)
	tri_dv, v.grad = v.grad.clone(), None
	tri_dk, k.grad = k.grad.clone(), None
	tri_dq, q.grad = q.grad.clone(), None
	# compare
	torch.testing.assert_close(ref_out, tri_out, atol=1e-2, rtol=0)

[whitneywhtsang]

@peterbell10 RuntimeError: "baddbmm_cuda" not implemented for 'Float8_e5m2' any suggestions on how to get reference output for fp8?