Add CausalMask support with new flash attention api

[ClarkChin08]

No description provided.

[petercad]

@ClarkChin08 -- You can derive SubgroupLayoutQK_ from TiledMMAQK_, so the user doesn't have to pass it in.

[petercad]

We should avoid making assumptions about the layout of MMA blocks within tSrS as they can easily broken by different choices for TiledMMAQK. Also, this code is assuming size<1>(tSrS) == 1.

Instead, I'd suggest using coordinate tensors to identify the coordinates of thread-owned data -- something like this (code below untested):

auto cS_thread = thr_mma_qk.partition_C(cP);     /* local S coordinates, within WG tile */
for (int i = 0; i < tSrS.size(); i++)
    if (get<1>(cS_thread(i)) >= seq_len - base_col)
        tSrS(i) = ElementS(-INFINITY);

Also, by using this method, you can avoid referring to SublayoutLayoutQK at all, because cS_thread is already aware of the subgroup tile's position within the workgroup tile.

[ClarkChin08]

Thanks @petercad , Let me have a try.

[Copilot]

Pull Request Overview

This PR adds support for causal masking in the flash attention implementation by introducing a new SubgroupLayoutQK template parameter and implementing the causal mask logic in the mainloop.

Key Changes:

• Added SubgroupLayoutQK template parameter to the collective mainloop and kernel interfaces
• Implemented causal masking logic that applies -INFINITY to attention scores beyond the causal boundary
• Updated the example runner to conditionally instantiate causal or non-causal configurations based on user options

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
applications/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp	Implements causal mask logic and removes the static assertion that previously blocked causal mask usage
applications/flash_attention_v2/kernel/xe_fhma_fwd_kernel.hpp	Adds subgroup layout type alias and computes sequence coordinates for causal masking
examples/06_bmg_flash_attention/xe_fmha_fwd_runner.hpp	Adds SubgroupLayoutQK template parameter to mainloop type
examples/06_bmg_flash_attention/06_xe_fmha_fwd.cpp	Conditionally selects causal or non-causal kernel based on is_causal option

[Copilot]

This calculation is overly complex and difficult to understand. Consider breaking it into intermediate variables with descriptive names to clarify the computation of tile offset within the subgroup layout.

Suggested change

	int seq_coord = cute::min(s.seq_len_qo, (blk_q * get<0>(TileShapeQK{}) + (sub_group_id / get<1>(shape(SubgroupLayoutQK{}))) * SGTileQ));
	// Break down the seq_coord calculation for clarity
	int tile_shape_qk_0 = get<0>(TileShapeQK{});
	int subgroup_layout_qk_1 = get<1>(shape(SubgroupLayoutQK{}));
	int blk_q_offset = blk_q * tile_shape_qk_0;
	int subgroup_tile_offset = (sub_group_id / subgroup_layout_qk_1) * SGTileQ;
	int raw_seq_coord = blk_q_offset + subgroup_tile_offset;
	int seq_coord = cute::min(s.seq_len_qo, raw_seq_coord);

[Copilot]

The causal mask condition logic lacks explanation. Add a comment describing what this inequality represents in terms of the attention matrix and why these specific offsets are used.

Suggested change

	int row_idx = seq_coord + m;
	int row_idx = seq_coord + m;
	// Causal mask: For each (row_idx, col_idx) in the attention matrix,
	// set positions where the query (row) would attend to future keys (col) to -INFINITY.
	// The offsets (full_tile_offset, discard_seq_coord) adjust for tiling and indexing,
	// ensuring that only positions where col_idx > row_idx (i.e., future positions)
	// are masked out, preserving causality in the attention mechanism.