# mypy: allow-untyped-defs """Triton Implementation of the flex_attention Kernel""" import logging import math from collections.abc import Sequence from dataclasses import dataclass from typing import Any, Optional, Union import sympy import torch from torch._inductor.utils import can_use_tma from torch._inductor.virtualized import V from ...ir import ComputedBuffer, ExternKernel, FixedLayout, TensorBox from ...lowering import empty, empty_strided, lowerings, register_lowering from ...select_algorithm import ( autotune_select_algorithm, SymbolicGridFn, TritonTemplate, ) from .common import ( build_subgraph_buffer, create_indices_fake, create_num_blocks_fake_generator, create_placeholder, get_fwd_subgraph_outputs, infer_dense_strides, load_template, maybe_realize, set_head_dim_values, SubgraphResults, ) from .flex_cpu import lower_cpu from .flex_decoding import _use_flex_decoding, create_flex_decoding_kernel log = logging.getLogger(__name__) aten = torch.ops.aten Expr = sympy.Expr @SymbolicGridFn def flex_attention_grid(batch_size, q_heads, num_queries, d_model, meta, *, cdiv): """How is this kernel parallelized? We create a grid of (ceil_div(n_queries, query_block_size), batch_size, num_heads) Each block is responsible for iterating over blocks of keys and values calculating the final attention output. """ return (cdiv(num_queries, meta["BLOCK_M"]), batch_size, q_heads) def get_float32_precision(): if ( ( torch.backends.cuda.matmul.fp32_precision == "ieee" if torch.backends.cuda.matmul.fp32_precision != "none" else torch.get_float32_matmul_precision() == "highest" ) or torch.version.hip or torch.mtia.is_available() ): return "'ieee'" else: return "'tf32'" flex_attention_template = TritonTemplate( name="flex_attention", grid=flex_attention_grid, source=load_template("flex_attention") + load_template("utilities") + load_template("common"), ) @register_lowering(torch.ops.higher_order.flex_attention, type_promotion_kind=None) def flex_attention( query, key, value, subgraph, block_mask, scale, kernel_options, score_mod_other_buffers, mask_mod_other_buffers, ): """The main lowering for the flex_attention hop This can currently lower to one of 3 templates: 1. Base Triton Template 2. Flex Decode Triton Template 3. Cpu specific CPP template """ if query.get_device().type == "cpu": return lower_cpu( query, key, value, subgraph, block_mask, scale, kernel_options, score_mod_other_buffers, mask_mod_other_buffers, ) # below is cuda path if device is not cpu # tl.dot does not support embedding size less than 16 small_dqk = V.graph.sizevars.evaluate_expr(sympy.Lt(query.get_size()[-1], 16)) small_dv = V.graph.sizevars.evaluate_expr(sympy.Lt(value.get_size()[-1], 16)) if small_dqk or small_dv: raise NotImplementedError( f"NYI: embedding dimension of the query, key, and value must be " f"at least 16 but got E={query.get_size()[-1]} and Ev={value.get_size()[-1]}" ) ( _, # q_length _, # kv_length kv_num_blocks, kv_indices, full_kv_num_blocks, full_kv_indices, q_num_blocks, q_indices, full_q_num_blocks, full_q_indices, SPARSE_Q_BLOCK_SIZE, SPARSE_KV_BLOCK_SIZE, mask_graph, ) = block_mask placeholder_inps = [ create_placeholder(name, dtype, query.get_device()) for name, dtype in [ ("score", query.get_dtype()), ("b", torch.int32), ("h", torch.int32), ("m", torch.int32), ("n", torch.int32), ] ] subgraph_buffer = build_subgraph_buffer( placeholder_inps + list(score_mod_other_buffers), subgraph ) mask_graph_placeholder_inps = [ create_placeholder(name, dtype, query.get_device()) for name, dtype in [ ("b", torch.int32), ("h", torch.int32), ("m", torch.int32), ("n", torch.int32), ] ] mask_graph_buffer = build_subgraph_buffer( mask_graph_placeholder_inps + list(mask_mod_other_buffers), mask_graph ) kernel_options = dict(kernel_options) # Mark symbols in custom kernel options as static shapes and add guards. kernel_options = { k: V.graph.sizevars.guard_int(v) if isinstance(v, sympy.Symbol) else v for k, v in kernel_options.items() } kernel_options.setdefault("FLOAT32_PRECISION", get_float32_precision()) enable_gqa = V.graph.sizevars.evaluate_expr( sympy.Ne(query.get_size()[1], key.get_size()[1]), ) if _use_flex_decoding(query, kv_indices, value, kernel_options, enable_gqa): return create_flex_decoding_kernel( query, key, value, block_mask, scale, kernel_options, subgraph_buffer, mask_graph_buffer, score_mod_other_buffers, mask_mod_other_buffers, ) ( query, key, value, kv_num_blocks, kv_indices, full_kv_num_blocks, full_kv_indices, q_num_blocks, q_indices, full_q_num_blocks, full_q_indices, ) = maybe_realize( [ query, key, value, kv_num_blocks, kv_indices, full_kv_num_blocks, full_kv_indices, q_num_blocks, q_indices, full_q_num_blocks, full_q_indices, ] ) score_mod_other_buffers = maybe_realize(score_mod_other_buffers) mask_mod_other_buffers = maybe_realize(mask_mod_other_buffers) Bq, Hq, seq_len_q, qk_head_dim = query.get_size() Bkv, Hkv, seq_len_kv, v_head_dim = value.get_size() assert V.graph.sizevars.evaluate_expr(sympy.Eq(Bq, Bkv) | sympy.Eq(Bkv, 1)), ( f"Bq and Bkv must broadcastable. Got Bq={Bq} and Bkv={Bkv}" ) assert V.graph.sizevars.evaluate_expr(sympy.Gt(seq_len_q, 0)), ( "Query length must be greater than 0" ) assert V.graph.sizevars.evaluate_expr(sympy.Gt(seq_len_kv, 0)), ( "Key length must be greater than 0" ) B = Bq if seq_len_q % 128 != 0 or seq_len_kv % 128 != 0: kernel_options.setdefault("IS_DIVISIBLE", False) else: kernel_options.setdefault("IS_DIVISIBLE", True) # NB it is okay that the v_head_dim is different # We are using these to match fill order of the output. q_strides = query.get_stride() # Construct output layout with strides matching the query. out_size = [B, Hq, seq_len_q, v_head_dim] out_strides = infer_dense_strides(out_size, q_strides) layout = FixedLayout( query.get_device(), query.get_dtype(), [B, Hq, seq_len_q, v_head_dim], stride=[sympy.sympify(s) for s in out_strides], ) # see NOTE:[TritonTemplates with multiple outputs] logsumexp_shape = [B, Hq, seq_len_q] logsumexp = empty_strided( logsumexp_shape, None, dtype=torch.float32, # The logsumexp is always stored in fp32 regardless of the input dtype device=query.get_device(), ) max_scores = empty_strided( logsumexp_shape, # Same shape as logsumexp None, dtype=torch.float32, # The max scores are always stored in fp32 regardless of the input dtype device=query.get_device(), ) kernel_options.setdefault("SM_SCALE", scale) # Determine GQA broadcast factor. gqa_shared_heads = Hq // Hkv kernel_options.setdefault("GQA_SHARED_HEADS", gqa_shared_heads) # Inside of Triton kernel, only apply partial masking if partial blocks are computed. # full_kv_num_blocks is None if partial blocks are not computed has_full_blocks = full_kv_num_blocks is not None kernel_options.setdefault("HAS_FULL_BLOCKS", has_full_blocks) if not has_full_blocks: full_kv_num_blocks, full_kv_indices = ( empty(0, device=query.get_device()) for _ in range(2) ) set_head_dim_values(kernel_options, qk_head_dim, v_head_dim, V.graph.sizevars) choices: list[Any] = [] dtype = query.get_dtype() head_dim = V.graph.sizevars.guard_int(query.get_size()[-1]) configs = V.choices.get_flex_attention_fwd_configs( head_dim, dtype, query.get_device().type ) # Mark SPARSE_KV_BLOCK_SIZE & SPARSE_Q_BLOCK_SIZE as static shapes and add guards. SPARSE_KV_BLOCK_SIZE = V.graph.sizevars.guard_int(SPARSE_KV_BLOCK_SIZE) SPARSE_Q_BLOCK_SIZE = V.graph.sizevars.guard_int(SPARSE_Q_BLOCK_SIZE) # Note, we don't need to pass in the captured buffers explicitly # because they're implicitly added by the score_mod function # We do need to explicitly pass it in for autotuning though. original_kernel_options = kernel_options.copy() # Default config for warp specialization num_consumer_groups, num_buffers_warp_spec = 0, 0 for conf in configs: cur_kernel_options = original_kernel_options.copy() # Performance tuning # Triton parameters # Remove prefix for forward kernels options and delete backward kernel options. for k in list(cur_kernel_options.keys()): if k.startswith("fwd_"): v = cur_kernel_options.pop(k) cur_kernel_options[k[4:]] = v if k.startswith("bwd_"): cur_kernel_options.pop(k) cur_kernel_options.setdefault("num_stages", conf.num_stages) cur_kernel_options.setdefault("num_warps", conf.num_warps) if cur_kernel_options.get("num_consumer_groups", False): cur_kernel_options.setdefault("num_consumer_groups", num_consumer_groups) cur_kernel_options.setdefault( "num_buffers_warp_spec", num_buffers_warp_spec ) # USE TMA = false by default cur_kernel_options.setdefault("USE_TMA", False) if cur_kernel_options["USE_TMA"] and can_use_tma(query, key, value): cur_kernel_options["USE_TMA"] = True cur_kernel_options.setdefault("BLOCK_M", conf.block_m) cur_kernel_options.setdefault("BLOCK_N", conf.block_n) # Blocksparse options cur_kernel_options.setdefault("SPARSE_Q_BLOCK_SIZE", SPARSE_Q_BLOCK_SIZE) cur_kernel_options.setdefault("SPARSE_KV_BLOCK_SIZE", SPARSE_KV_BLOCK_SIZE) if ( cur_kernel_options["SPARSE_KV_BLOCK_SIZE"] % cur_kernel_options["BLOCK_N"] != 0 or cur_kernel_options["SPARSE_Q_BLOCK_SIZE"] % cur_kernel_options["BLOCK_M"] != 0 ): if len(configs) == 1: raise ValueError( f"Q and KV block size must be divisible by BLOCK_M and BLOCK_N. We " f"got Q_BLOCK_SIZE={cur_kernel_options['SPARSE_Q_BLOCK_SIZE']} and " f"KV_BLOCK_SIZE={cur_kernel_options['SPARSE_KV_BLOCK_SIZE']}." ) continue # ROCm specific kernargs for attrib in ["kpack", "matrix_instr_nonkdim", "waves_per_eu"]: if hasattr(conf, attrib): cur_kernel_options[attrib] = getattr(conf, attrib) error = flex_attention_template.maybe_append_choice( choices=choices, input_nodes=[ query, key, value, logsumexp, max_scores, kv_num_blocks, kv_indices, full_kv_num_blocks, full_kv_indices, ], layout=layout, subgraphs=[ subgraph_buffer, mask_graph_buffer, ], mutated_inputs=[ logsumexp, max_scores, ], call_sizes=query.get_size(), **cur_kernel_options, ) if error is not None and len(configs) == 1: raise error inputs_for_autotuning = ( [ query, key, value, logsumexp, max_scores, kv_num_blocks, kv_indices, full_kv_num_blocks, full_kv_indices, ] + list(score_mod_other_buffers) + list(mask_mod_other_buffers) ) input_gen_fns = { 5: create_num_blocks_fake_generator(kv_indices), 6: create_indices_fake, 7: create_num_blocks_fake_generator(full_kv_indices), 8: create_indices_fake, } out = autotune_select_algorithm( "flex_attention", choices, # Need to filter out symbols since there is an invariant # that all input_nodes are of type IRNode [x for x in inputs_for_autotuning if isinstance(x, torch._inductor.ir.IRNode)], layout, input_gen_fns=input_gen_fns, ) # need subgraph inputs and outputs to analyze all symints used in flex attention out.data.data.subgraph_inps = list(score_mod_other_buffers) + list( mask_mod_other_buffers ) out.data.data.subgraph_outs = get_fwd_subgraph_outputs( subgraph_buffer, mask_graph_buffer ) return (out, logsumexp, max_scores) # ---------------------------- Backward HOP Implementation ---------------------------- @SymbolicGridFn def flex_attention_backward_grid( batch_size, q_heads, num_queries, d_model, kv_heads, num_key_value, meta, *, cdiv ): """How is this kernel parallelized? We create a grid of (ceil_div(n_queries, query_block_size) * heads_ratio + ceil_div(n_kv, kv_block_size), batch_size, kv_heads) Currently this is only parallelizing over batch* kv_heads, but we can, and want to parallelize over ceil_div(q_heads//kv_heads * num_key_value, key_value_block_size). To do this will either require atomic updates to some grad values or to have a two pass kernel design. """ return ( cdiv(num_queries, meta["BLOCK_M2"]) * (q_heads // kv_heads) + cdiv(num_key_value, meta["BLOCK_N1"]), batch_size, kv_heads, ) flex_attention_backward_template = TritonTemplate( name="flex_attention_backward", grid=flex_attention_backward_grid, source=load_template("flex_backwards") + load_template("utilities"), ) def validate_joint_graph(joint_graph: torch.fx.Graph): """We do some pre lowering graph checks in order to raise nicer error messages""" for node in joint_graph.nodes: if ( node.op == "call_function" and node.target == torch.ops.flex_lib.zeros_and_scatter.default ): for user in node.users: if user.op != "output": raise NotImplementedError( "Using multiple indexing operations on the same tensor that requires gradients " "in a score_mod function is not currently supported. " "This typically happens when indexing the same tensor multiple times, like:\n\n" " def score_mod(score, b, h, q_idx, kv_idx):\n" " return score + bias[q_idx] + bias[kv_idx] # bias used twice!\n\n" "A valid workaround is to clone() the tensors that will be indexed multiple times. For example:\n\n" " bias1 = bias.clone()\n" " def score_mod(score, b, h, q_idx, kv_idx):\n" " return score + bias[q_idx] + bias1[kv_idx]\n\n" "Note that this solution will use additional memory." ) return @dataclass(frozen=True) class JointOutputResult: """Results from processing joint outputs.""" grad_input: ComputedBuffer captured_grads_compute: list[ComputedBuffer] captured_grads: list[Optional[TensorBox]] mutated_grads: list[TensorBox] def process_joint_outputs( all_joint_outputs: SubgraphResults, num_placeholders: int ) -> JointOutputResult: """Process joint outputs and extract various buffers needed for lowering Args: all_joint_outputs: List of all the outputs from build_subgraphs num_placeholders: The number of placeholder inputs, used to skip over unused backward compute buffers Returns: JointOutputResult containing processed buffers and gradients """ assert isinstance(all_joint_outputs, list) assert all_joint_outputs[0] is not None, ( "joint_subgraph_buffer is None - this is a bug!" ) joint_buffer = all_joint_outputs[0] other_grads = all_joint_outputs[num_placeholders - 1 :] # outer_grads has the structure: Len(other_buffer_grads) if buffer doesn't require grad than it will be None # We only grab the buffers that require grad for inlining into kernel grads_compute = [buf for buf in other_grads if buf is not None] def get_out(buf): if buf is None: return None assert isinstance(buf, ComputedBuffer) assert buf.name is not None return TensorBox.create(V.graph.get_buffer(buf.name)) grads_out = [get_out(x) for x in other_grads] mutated_grads = [buf for buf in grads_out if buf is not None] return JointOutputResult( grad_input=joint_buffer, captured_grads_compute=grads_compute, captured_grads=grads_out, mutated_grads=mutated_grads, ) # TODO: We probably also need a layout constraint? @register_lowering( torch.ops.higher_order.flex_attention_backward, type_promotion_kind=None ) def flex_attention_backward(*args, **kwargs): """Lowering for the flex_attention_backward op in triton""" ( query, key, value, out, logsumexp, grad_out, grad_logsumexp, fw_graph, joint_graph, block_mask, scale, kernel_options, score_mod_other_buffers, mask_mod_other_buffers, ) = args ( _, # q_length _, # kv_length kv_num_blocks, kv_indices, full_kv_num_blocks, full_kv_indices, q_num_blocks, q_indices, full_q_num_blocks, full_q_indices, SPARSE_Q_BLOCK_SIZE, SPARSE_KV_BLOCK_SIZE, mask_graph, ) = block_mask ( query, key, value, grad_out, kv_num_blocks, kv_indices, full_kv_num_blocks, full_kv_indices, q_num_blocks, q_indices, full_q_num_blocks, full_q_indices, ) = maybe_realize( [ query, key, value, grad_out, kv_num_blocks, kv_indices, full_kv_num_blocks, full_kv_indices, q_num_blocks, q_indices, full_q_num_blocks, full_q_indices, ] ) device = query.get_device() dtype = query.get_dtype() Bq, Hq, seq_len_q, qk_head_dim = query.get_size() Bkv, Hkv, seq_len_kv, v_head_dim = value.get_size() assert V.graph.sizevars.evaluate_expr(sympy.Eq(Bq, Bkv) | sympy.Eq(Bkv, 1)), ( f"Bq and Bkv must broadcastable. Got Bq={Bq} and Bkv={Bkv}" ) kernel_options = dict(kernel_options) # Mark symbols in custom kernel options as static shapes and add guards. kernel_options = { k: V.graph.sizevars.guard_int(v) if isinstance(v, sympy.Symbol) else v for k, v in kernel_options.items() } kernel_options.setdefault("FLOAT32_PRECISION", get_float32_precision()) seq_q_divisible = V.graph.sizevars.statically_known_true(seq_len_q % 128 == 0) seq_kv_divisible = V.graph.sizevars.statically_known_true(seq_len_kv % 128 == 0) if seq_q_divisible and seq_kv_divisible: kernel_options.setdefault("IS_DIVISIBLE", True) else: kernel_options.setdefault("IS_DIVISIBLE", False) fwd_placeholder_inps = [ create_placeholder(name, dtype, device) for name, dtype in [ ("score", dtype), ("b", torch.int32), ("h", torch.int32), ("m", torch.int32), ("n", torch.int32), ] ] fw_subgraph_buffer = build_subgraph_buffer( fwd_placeholder_inps + list(score_mod_other_buffers), fw_graph ) joint_placeholder_inps = fwd_placeholder_inps + [ create_placeholder("grad_score_mod", dtype, device) ] # Sometimes we have weird unused nodes here joint_graph.graph_module.graph.eliminate_dead_code() # It is hard to raise nice errors for some joint graphs during subgraph lowering # This lets us do some checks before attempting to lower validate_joint_graph(joint_graph.graph_module.graph) all_joint_outputs = build_subgraph_buffer( joint_placeholder_inps + list(score_mod_other_buffers), joint_graph, ) joint_outputs = process_joint_outputs( all_joint_outputs, len(joint_placeholder_inps) ) mask_graph_placeholder_inps = [ create_placeholder(name, dtype, query.get_device()) for name, dtype in [ ("b", torch.int32), ("h", torch.int32), ("m", torch.int32), ("n", torch.int32), ] ] mask_graph_buffer = build_subgraph_buffer( mask_graph_placeholder_inps + list(mask_mod_other_buffers), mask_graph ) mask_graph_buffer = mask_graph_buffer # Construct layout with stride order matching K key_size = [Bq, Hkv, seq_len_kv, qk_head_dim] key_strides = infer_dense_strides(key_size, key.get_stride()) layout_broadcasted_k = FixedLayout( key.get_device(), key.get_dtype(), key_size, stride=[sympy.sympify(s) for s in key_strides], ) # Create delta which will is needed for the bwd's kernel grad_lse_exp2 = lowerings[aten.mul](grad_logsumexp, 1 / math.log(2)) mul_delta = lowerings[aten.mul](out, grad_out) delta = lowerings[aten.sum](mul_delta, axis=-1) delta = lowerings[aten.sub](delta, grad_lse_exp2) delta = ExternKernel.require_contiguous(delta) grad_lse_exp2, delta = maybe_realize([grad_lse_exp2, delta]) # # see NOTE:[TritonTemplates with multiple outputs] query_size = [Bq, Hq, seq_len_q, qk_head_dim] grad_query_strides = infer_dense_strides(query_size, query.get_stride()) grad_query = empty_strided( query_size, stride=[sympy.sympify(s) for s in grad_query_strides], dtype=query.get_dtype(), device=query.get_device(), ) # Construct output layout with stride order matching value value_size = [Bq, Hkv, seq_len_kv, v_head_dim] value_strides = infer_dense_strides(value_size, value.get_stride()) broadcasted_grad_value = empty_strided( value_size, stride=[sympy.sympify(s) for s in value_strides], dtype=value.get_dtype(), device=value.get_device(), ) kernel_options.setdefault("SM_SCALE", scale) # Determine GQA factor gqa_shared_heads = Hq // Hkv kernel_options.setdefault("GQA_SHARED_HEADS", gqa_shared_heads) # Inside of Triton kernel, only apply partial masking if partial blocks are computed. # full_kv_num_blocks is torch.zeros([1, 1, 1]) if partial blocks are not computed. has_full_blocks = full_kv_num_blocks is not None kernel_options.setdefault("HAS_FULL_BLOCKS", has_full_blocks) if not has_full_blocks: full_kv_num_blocks, full_kv_indices, full_q_num_blocks, full_q_indices = ( empty(0, device=query.get_device()) for _ in range(4) ) set_head_dim_values(kernel_options, qk_head_dim, v_head_dim, V.graph.sizevars) SPARSE_Q_BLOCK_SIZE = V.graph.sizevars.guard_int(SPARSE_Q_BLOCK_SIZE) SPARSE_KV_BLOCK_SIZE = V.graph.sizevars.guard_int(SPARSE_KV_BLOCK_SIZE) choices: list[Any] = [] dtype = query.get_dtype() head_dim = V.graph.sizevars.guard_int(query.get_size()[-1]) configs = V.choices.get_flex_attention_bwd_configs( head_dim, dtype, query.get_device().type ) # Default config for warp specialization num_consumer_groups, num_buffers_warp_spec = 0, 0 original_kernel_options = kernel_options.copy() for conf in configs: if ( SPARSE_KV_BLOCK_SIZE % conf.block_m != 0 or SPARSE_Q_BLOCK_SIZE % conf.block_m != 0 or SPARSE_KV_BLOCK_SIZE % conf.block_n != 0 or SPARSE_Q_BLOCK_SIZE % conf.block_n != 0 ): continue # Performance tuning # Triton heuristics cur_kernel_options = original_kernel_options.copy() # Remove prefix for backward kernels options and delete forward kernel options. for k in list(cur_kernel_options.keys()): if k.startswith("bwd_"): v = cur_kernel_options.pop(k) cur_kernel_options[k[4:]] = v if k.startswith("fwd_"): cur_kernel_options.pop(k) cur_kernel_options.setdefault("num_warps", conf.num_warps) cur_kernel_options.setdefault("num_stages", conf.num_stages) if cur_kernel_options.get("num_consumer_groups", False): cur_kernel_options.setdefault("num_consumer_groups", num_consumer_groups) cur_kernel_options.setdefault( "num_buffers_warp_spec", num_buffers_warp_spec ) cur_kernel_options.setdefault("BLOCK_M1", conf.block_m) cur_kernel_options.setdefault("BLOCK_N1", conf.block_n) cur_kernel_options.setdefault("BLOCK_M2", conf.block_n) cur_kernel_options.setdefault("BLOCK_N2", conf.block_m) # Blocksparse options cur_kernel_options.setdefault("SPARSE_Q_BLOCK_SIZE", SPARSE_Q_BLOCK_SIZE) cur_kernel_options.setdefault("SPARSE_KV_BLOCK_SIZE", SPARSE_KV_BLOCK_SIZE) # ROCm specific kernargs for attrib in ["kpack", "matrix_instr_nonkdim", "waves_per_eu"]: if hasattr(conf, attrib): cur_kernel_options[attrib] = getattr(conf, attrib) flex_attention_backward_template.maybe_append_choice( choices=choices, input_nodes=[ query, key, value, logsumexp, delta, grad_out, grad_query, broadcasted_grad_value, kv_num_blocks, kv_indices, q_num_blocks, q_indices, full_kv_num_blocks, full_kv_indices, full_q_num_blocks, full_q_indices, ], layout=layout_broadcasted_k, # We use store_output only for grad_key subgraphs=[ fw_subgraph_buffer, joint_outputs.grad_input, mask_graph_buffer, joint_outputs.captured_grads_compute, ], mutated_inputs=[ grad_query, broadcasted_grad_value, *joint_outputs.mutated_grads, ], call_sizes=query.get_size() + key.get_size()[1:3], **cur_kernel_options, ) inputs_for_autotuning = ( [ query, key, value, logsumexp, delta, grad_out, grad_query, broadcasted_grad_value, kv_num_blocks, kv_indices, q_num_blocks, q_indices, full_kv_num_blocks, full_kv_indices, full_q_num_blocks, full_q_indices, ] + list(score_mod_other_buffers) + list(mask_mod_other_buffers) + joint_outputs.mutated_grads ) input_gen_fns = { 8: create_num_blocks_fake_generator(kv_indices), # kv_num_blocks 9: create_indices_fake, 10: create_num_blocks_fake_generator(q_indices), # q_num_blocks 11: create_indices_fake, 12: create_num_blocks_fake_generator(full_kv_indices), # full_kv_num_blocks 13: create_indices_fake, 14: create_num_blocks_fake_generator(full_q_indices), # full_q_num_blocks 15: create_indices_fake, } broadcasted_grad_key = autotune_select_algorithm( "flex_attention_backward", choices, [x for x in inputs_for_autotuning if isinstance(x, torch._inductor.ir.IRNode)], layout_broadcasted_k, input_gen_fns=input_gen_fns, ) # [Bq, Hkv, seq_len_kv, k_head_dim] # need subgraph inputs and outputs to analyze all symints used in flex attention broadcasted_grad_key.data.data.subgraph_inps = list(score_mod_other_buffers) + list( mask_mod_other_buffers ) broadcasted_grad_key.data.data.subgraph_outs = get_bwd_subgraph_outputs( fw_subgraph_buffer, mask_graph_buffer, joint_outputs ) if V.graph.sizevars.evaluate_expr(sympy.Eq(Bq, Bkv)): grad_key = broadcasted_grad_key grad_value = broadcasted_grad_value else: assert V.graph.sizevars.evaluate_expr(sympy.Gt(Bq, 1) & sympy.Eq(Bkv, 1)), ( f"Bq and Bkv must broadcastable. " f"Got Bq={V.graph.sizevars.evaluate_expr(Bq)} " f"and Bkv={V.graph.sizevars.evaluate_expr(Bkv)}" ) grad_key = lowerings[aten.sum](broadcasted_grad_key, axis=0, keepdims=True) grad_value = lowerings[aten.sum](broadcasted_grad_value, axis=0, keepdims=True) return (grad_query, grad_key, grad_value, tuple(joint_outputs.captured_grads)) def get_bwd_subgraph_outputs( subgraph_buffer: SubgraphResults, mask_graph_buffer: SubgraphResults, joint_outputs: JointOutputResult, ) -> list[Optional[Union[ComputedBuffer, TensorBox]]]: subgraph_buffer = ( subgraph_buffer if isinstance(subgraph_buffer, Sequence) else [subgraph_buffer] ) mask_graph_buffer = ( mask_graph_buffer if isinstance(mask_graph_buffer, Sequence) else [mask_graph_buffer] ) joint_output_buffers = [ joint_outputs.grad_input, *joint_outputs.captured_grads_compute, *joint_outputs.captured_grads, *joint_outputs.mutated_grads, ] return [*subgraph_buffer, *mask_graph_buffer, *joint_output_buffers]