
{{def_kernel("Q", "K", "V", "LOGSUMEXP")}}
    from flash_attn.cute.interface import _flash_attn_fwd

    # Transpose tensors for _flash_attn_fwd compatibility (B,H,M,D) -> (B,M,H,D)
    q_transposed = Q.transpose(1, 2)
    k_transposed = K.transpose(1, 2)
    v_transposed = V.transpose(1, 2)

    @cute.jit
    def score_mod(tSrS_ssa, b_idx, h_idx, q_idx, kv_idx, buffers):
        {{unpack_buffers()}}
        {{ modification(
            subgraph_number=0,
            output_name="tSrS_ssa",
            score="tSrS_ssa",
            b="b_idx",
            h="h_idx",
            m="q_idx",
            n="kv_idx",
            out="tSrS_ssa"
        ) | indent_except_first(2) }}
        return tSrS_ssa

    # (B,M,H,D) -> (B,H,M,D)
    output = {{get_output()}}
    output_transposed = output.transpose(1, 2)

    # Collect any additional tensor buffers that were added during modifications
    {% set tensor_buffers = get_tensor_buffers() -%}
    {% if tensor_buffers -%}
    buffers = [{% for buffer in tensor_buffers %}{{buffer}}{% if not loop.last %}, {% endif %}{% endfor %}]
    buffers = list(buffers)
    {% else -%}
    buffers = None
    {% endif -%}

    # Out and LSE filled inplace
    _flash_attn_fwd(
        q_transposed,
        k_transposed,
        v_transposed,
        softmax_scale={{SM_SCALE}},
        causal={{CAUSAL}},
        return_lse=True,
        score_mod=score_mod,
        out=output_transposed,
        lse=LOGSUMEXP,
        buffers=buffers
    )