| import torch |
| import triton |
| import triton.language as tl |
|
|
|
|
|
|
| def get_autotuning_configs(q_tile_sizes: list): |
| """Generate autotuning configurations optimized for H100.""" |
| warps = [4, 8] |
| stages = [2, 3] |
|
|
| return [ |
| triton.Config({'q_tile_size': t}, num_warps=w, num_stages=s) |
| for t in q_tile_sizes |
| for w in warps |
| for s in stages |
| ] |
|
|
|
|
| @triton.autotune( |
| configs=get_autotuning_configs([64, 128]), |
| key=['seq_len', 'feature_dim'], |
| ) |
| @triton.jit |
| def mosaic_attn_fwd_kernel( |
| q_ptr, k_ptr, v_ptr, output_ptr, lse_ptr, block_indices_ptr, |
| softmax_scale: tl.constexpr, |
| seq_len: tl.constexpr, |
| num_kv_heads: tl.constexpr, |
| num_q_heads: tl.constexpr, |
| q_heads_per_kv_head: tl.constexpr, |
| feature_dim: tl.constexpr, |
| kv_block_size: tl.constexpr, |
| num_kv_blocks_per_q_block: tl.constexpr, |
| q_tile_size: tl.constexpr, |
| ): |
| """ |
| Sparse attention forward kernel: |
| for each query tile (i.e. block chunk), for each query head, attend to a subset of key/value blocks. |
| """ |
| LOG2_E: tl.constexpr = 1.44269504089 |
|
|
| q_tile_id = tl.program_id(0) |
| q_head_id = tl.program_id(1) |
| batch_kv_head_id = tl.program_id(2) |
|
|
| batch_idx = batch_kv_head_id // num_kv_heads |
| kv_head_idx = batch_kv_head_id % num_kv_heads |
| q_head_idx = kv_head_idx * q_heads_per_kv_head + q_head_id |
|
|
| batch_offset = batch_idx * seq_len |
| q_tile_start = q_tile_id * q_tile_size |
| num_blocks_in_seq = seq_len // kv_block_size |
| tiles_per_block = kv_block_size // q_tile_size |
| q_block_id = q_tile_id // tiles_per_block |
|
|
| block_indices_offset = ( |
| batch_idx * num_blocks_in_seq * num_kv_heads * num_kv_blocks_per_q_block + |
| q_block_id * num_kv_heads * num_kv_blocks_per_q_block + |
| kv_head_idx * num_kv_blocks_per_q_block |
| ) |
|
|
| q_base_ptr = q_ptr + batch_offset * num_q_heads * feature_dim + q_head_idx * feature_dim |
| k_base_ptr = k_ptr + batch_offset * num_kv_heads * feature_dim + kv_head_idx * feature_dim |
| v_base_ptr = v_ptr + batch_offset * num_kv_heads * feature_dim + kv_head_idx * feature_dim |
|
|
| q_tile_ptr = tl.make_block_ptr( |
| base=q_base_ptr, |
| shape=(seq_len, feature_dim), |
| strides=(num_q_heads * feature_dim, 1), |
| offsets=(q_tile_start, 0), |
| block_shape=(q_tile_size, feature_dim), |
| order=(1, 0) |
| ) |
|
|
| output_tile_ptr = tl.make_block_ptr( |
| base=output_ptr + batch_offset * num_q_heads * feature_dim + q_head_idx * feature_dim, |
| shape=(seq_len, feature_dim), |
| strides=(num_q_heads * feature_dim, 1), |
| offsets=(q_tile_start, 0), |
| block_shape=(q_tile_size, feature_dim), |
| order=(1, 0) |
| ) |
|
|
| lse_base_ptr = lse_ptr + (batch_offset + q_tile_start) * num_q_heads + tl.arange(0, q_tile_size) * num_q_heads + q_head_idx |
|
|
| output_accum = tl.zeros([q_tile_size, feature_dim], dtype=tl.float32) |
| max_scores = tl.full([q_tile_size], float('-inf'), dtype=tl.float32) |
| sum_exp_scores = tl.zeros([q_tile_size], dtype=tl.float32) |
|
|
| q_tile = tl.load(q_tile_ptr) |
| q_tile = (q_tile * softmax_scale * LOG2_E).to(tl.bfloat16) |
|
|
| for i in range(num_kv_blocks_per_q_block): |
| kv_block_start = kv_block_size * tl.load(block_indices_ptr + block_indices_offset + i).to(tl.int32) |
|
|
| k_block_ptr = tl.make_block_ptr( |
| base=k_base_ptr, |
| shape=(feature_dim, seq_len), |
| strides=(1, num_kv_heads * feature_dim), |
| offsets=(0, kv_block_start), |
| block_shape=(feature_dim, kv_block_size), |
| order=(1, 0) |
| ) |
|
|
| v_block_ptr = tl.make_block_ptr( |
| base=v_base_ptr, |
| shape=(seq_len, feature_dim), |
| strides=(num_kv_heads * feature_dim, 1), |
| offsets=(kv_block_start, 0), |
| block_shape=(kv_block_size, feature_dim), |
| order=(1, 0) |
| ) |
|
|
| k_block = tl.load(k_block_ptr).to(tl.bfloat16) |
| v_block = tl.load(v_block_ptr).to(tl.bfloat16) |
|
|
| attention_scores = tl.dot(q_tile, k_block) |
|
|
| new_max = tl.max(attention_scores, axis=1) |
| old_max = max_scores |
| max_scores = tl.maximum(max_scores, new_max) |
| rescale = tl.exp2(old_max - max_scores) |
| attention_probs = tl.exp2(attention_scores - max_scores[:, None]) |
| sum_exp_scores = sum_exp_scores * rescale + tl.sum(attention_probs, axis=1) |
|
|
| output_accum = output_accum * rescale[:, None] |
| output_accum += tl.dot(attention_probs.to(tl.bfloat16), v_block) |
|
|
| final_output = output_accum / sum_exp_scores[:, None] |
| log_sum_exp = (max_scores + tl.log2(sum_exp_scores)) |
|
|
| tl.store(output_tile_ptr, final_output.to(q_ptr.dtype.element_ty)) |
| tl.store(lse_base_ptr, log_sum_exp.to(tl.float32)) |
|
|
|
|
| def mosaic_attn_fwd( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| block_indices: torch.LongTensor, |
| block_size: int, |
| softmax_scale: float, |
| ): |
| batch_size, seq_len, num_kv_heads, feature_dim = k.shape |
| num_q_heads = q.shape[2] |
| num_kv_blocks_per_q_block = block_indices.shape[-1] |
| q_heads_per_kv_head = num_q_heads // num_kv_heads |
|
|
| output = torch.empty(batch_size, seq_len, num_q_heads, feature_dim, dtype=v.dtype, device=q.device) |
| lse = torch.empty(batch_size, seq_len, num_q_heads, dtype=torch.float32, device=q.device) |
|
|
| grid = lambda META: ( |
| triton.cdiv(seq_len, META['q_tile_size']), |
| q_heads_per_kv_head, |
| batch_size * num_kv_heads |
| ) |
|
|
| mosaic_attn_fwd_kernel[grid]( |
| q_ptr = q, |
| k_ptr = k, |
| v_ptr = v, |
| output_ptr = output, |
| lse_ptr = lse, |
| block_indices_ptr = block_indices, |
| softmax_scale = softmax_scale, |
| seq_len = seq_len, |
| num_kv_heads = num_kv_heads, |
| num_q_heads = num_q_heads, |
| q_heads_per_kv_head = q_heads_per_kv_head, |
| feature_dim = feature_dim, |
| kv_block_size = block_size, |
| num_kv_blocks_per_q_block = num_kv_blocks_per_q_block, |
| ) |
|
|
| return output, lse |
|
|
|
|
| @triton.autotune( |
| configs=get_autotuning_configs([64, 128]), |
| key=['seq_len', 'feature_dim'], |
| ) |
| @triton.jit |
| def mosaic_attn_bwd_q_kernel( |
| q_ptr, k_ptr, v_ptr, lse_ptr, delta_ptr, grad_o_ptr, grad_q_ptr, block_indices_ptr, |
| softmax_scale: tl.constexpr, |
| seq_len: tl.constexpr, |
| num_kv_heads: tl.constexpr, |
| num_q_heads: tl.constexpr, |
| q_heads_per_kv_head: tl.constexpr, |
| feature_dim: tl.constexpr, |
| kv_block_size: tl.constexpr, |
| num_kv_blocks_per_q_block: tl.constexpr, |
| q_tile_size: tl.constexpr, |
| ): |
| LOG2_E: tl.constexpr = 1.44269504089 |
| LN_2: tl.constexpr = 0.69314718056 |
|
|
| q_tile_id = tl.program_id(0) |
| q_head_id = tl.program_id(1) |
| batch_kv_head_id = tl.program_id(2) |
|
|
| batch_idx = batch_kv_head_id // num_kv_heads |
| kv_head_idx = batch_kv_head_id % num_kv_heads |
| q_head_idx = kv_head_idx * q_heads_per_kv_head + q_head_id |
|
|
| batch_offset = batch_idx * seq_len |
| q_tile_start = q_tile_id * q_tile_size |
| tiles_per_block = kv_block_size // q_tile_size |
| q_block_id = q_tile_id // tiles_per_block |
| num_q_blocks = seq_len // kv_block_size |
|
|
| block_indices_offset = ( |
| batch_idx * num_q_blocks * num_kv_heads * num_kv_blocks_per_q_block + |
| q_block_id * num_kv_heads * num_kv_blocks_per_q_block + |
| kv_head_idx * num_kv_blocks_per_q_block |
| ) |
|
|
| q_offsets = ( |
| tl.arange(0, q_tile_size)[:, None] * num_q_heads * feature_dim + |
| q_head_idx * feature_dim + |
| tl.arange(0, feature_dim)[None, :] |
| ) |
|
|
| lse_offsets = tl.arange(0, q_tile_size) * num_q_heads + q_head_idx |
|
|
| q_base_ptr = q_ptr + (batch_offset + q_tile_start) * num_q_heads * feature_dim |
| grad_o_base_ptr = grad_o_ptr + (batch_offset + q_tile_start) * num_q_heads * feature_dim |
| delta_base_ptr = delta_ptr + (batch_offset + q_tile_start) * num_q_heads |
| lse_base_ptr = lse_ptr + (batch_offset + q_tile_start) * num_q_heads |
| grad_q_base_ptr = grad_q_ptr + (batch_offset + q_tile_start) * num_q_heads * feature_dim |
|
|
| grad_q_accum = tl.zeros([q_tile_size, feature_dim], dtype=tl.float32) |
|
|
| q_tile = tl.load(q_base_ptr + q_offsets) |
| q_tile = (q_tile * softmax_scale * LOG2_E).to(tl.bfloat16) |
|
|
| grad_o_tile = tl.load(grad_o_base_ptr + q_offsets).to(tl.bfloat16) |
| delta_vals = tl.load(delta_base_ptr + lse_offsets) |
| lse_vals = tl.load(lse_base_ptr + lse_offsets).to(tl.float32) |
|
|
| for i in range(num_kv_blocks_per_q_block): |
| kv_block_idx = tl.load(block_indices_ptr + block_indices_offset + i).to(tl.int32) |
|
|
| k_block_ptr = tl.make_block_ptr( |
| base=k_ptr + (batch_offset * num_kv_heads + kv_head_idx) * feature_dim, |
| shape=(feature_dim, seq_len), |
| strides=(1, num_kv_heads * feature_dim), |
| offsets=(0, kv_block_idx * kv_block_size), |
| block_shape=(feature_dim, kv_block_size), |
| order=(0, 1) |
| ) |
|
|
| v_block_ptr = tl.make_block_ptr( |
| base=v_ptr + (batch_offset * num_kv_heads + kv_head_idx) * feature_dim, |
| shape=(feature_dim, seq_len), |
| strides=(1, num_kv_heads * feature_dim), |
| offsets=(0, kv_block_idx * kv_block_size), |
| block_shape=(feature_dim, kv_block_size), |
| order=(0, 1) |
| ) |
|
|
| k_block = tl.load(k_block_ptr).to(tl.bfloat16) |
| v_block = tl.load(v_block_ptr).to(tl.bfloat16) |
|
|
| attention_scores = tl.dot(q_tile, k_block) |
| attention_probs = tl.exp2(attention_scores - lse_vals[:, None]) * LN_2 |
|
|
| grad_times_v = tl.dot(grad_o_tile, v_block) |
| grad_scores = attention_probs * (grad_times_v - delta_vals[:, None]) |
| grad_q_accum += tl.dot(grad_scores.to(tl.bfloat16), tl.trans(k_block.to(tl.bfloat16))) |
|
|
| grad_q_accum = grad_q_accum * softmax_scale * LOG2_E |
| tl.store(grad_q_base_ptr + q_offsets, grad_q_accum.to(q_ptr.dtype.element_ty)) |
|
|
|
|
| @torch.compile |
| @torch.no_grad() |
| def mosaic_block_mask( |
| block_indices: torch.LongTensor, |
| ): |
| batch_size, num_blocks, num_heads, _ = block_indices.shape |
|
|
| block_mask = torch.zeros( |
| batch_size, num_blocks, num_heads, num_blocks, |
| dtype=torch.bool, device=block_indices.device |
| ) |
|
|
| batch_idx = torch.arange(batch_size, device=block_indices.device)[:, None, None, None] |
| q_block_idx = torch.arange(num_blocks, device=block_indices.device)[None, :, None, None] |
| head_idx = torch.arange(num_heads, device=block_indices.device)[None, None, :, None] |
|
|
| block_mask[batch_idx, q_block_idx, head_idx, block_indices] = True |
|
|
| block_mask_transposed = block_mask.permute(0, 2, 3, 1).contiguous() |
|
|
| return block_mask_transposed |
|
|
|
|
| @triton.autotune( |
| configs=get_autotuning_configs([16, 32]), |
| key=['seq_len', 'feature_dim'], |
| ) |
| @triton.jit |
| def mosaic_attn_bwd_kv_kernel( |
| q_ptr, k_ptr, v_ptr, lse_ptr, delta_ptr, |
| grad_o_ptr, grad_k_ptr, grad_v_ptr, |
| block_mask_ptr, |
| softmax_scale: tl.constexpr, |
| seq_len: tl.constexpr, |
| num_kv_heads: tl.constexpr, |
| num_q_heads: tl.constexpr, |
| q_heads_per_kv_head: tl.constexpr, |
| feature_dim: tl.constexpr, |
| kv_block_size: tl.constexpr, |
| q_tile_size: tl.constexpr, |
| ): |
| LOG2_E: tl.constexpr = 1.44269504089 |
| LN_2: tl.constexpr = 0.69314718056 |
|
|
| kv_block_id = tl.program_id(0) |
| batch_kv_head_id = tl.program_id(1) |
|
|
| batch_idx = batch_kv_head_id // num_kv_heads |
| kv_head_idx = batch_kv_head_id % num_kv_heads |
| batch_offset = batch_idx * seq_len |
|
|
| num_blocks_in_seq = seq_len // kv_block_size |
| tiles_per_block = kv_block_size // q_tile_size |
|
|
| fine_mask_start = ( |
| batch_idx * num_kv_heads * num_blocks_in_seq * num_blocks_in_seq + |
| kv_head_idx * num_blocks_in_seq * num_blocks_in_seq + |
| kv_block_id * num_blocks_in_seq |
| ) |
|
|
| k_block_ptr = tl.make_block_ptr( |
| k_ptr + (batch_offset * num_kv_heads + kv_head_idx) * feature_dim, |
| (seq_len, feature_dim), (num_kv_heads * feature_dim, 1), |
| (kv_block_id * kv_block_size, 0), (kv_block_size, feature_dim), (1, 0) |
| ) |
|
|
| v_block_ptr = tl.make_block_ptr( |
| v_ptr + (batch_offset * num_kv_heads + kv_head_idx) * feature_dim, |
| (seq_len, feature_dim), (num_kv_heads * feature_dim, 1), |
| (kv_block_id * kv_block_size, 0), (kv_block_size, feature_dim), (1, 0) |
| ) |
|
|
| grad_k_ptr = tl.make_block_ptr( |
| grad_k_ptr + (batch_offset * num_kv_heads + kv_head_idx) * feature_dim, |
| (seq_len, feature_dim), (num_kv_heads * feature_dim, 1), |
| (kv_block_id * kv_block_size, 0), (kv_block_size, feature_dim), (1, 0) |
| ) |
|
|
| grad_v_ptr = tl.make_block_ptr( |
| grad_v_ptr + (batch_offset * num_kv_heads + kv_head_idx) * feature_dim, |
| (seq_len, feature_dim), (num_kv_heads * feature_dim, 1), |
| (kv_block_id * kv_block_size, 0), (kv_block_size, feature_dim), (1, 0) |
| ) |
|
|
| k_block = tl.load(k_block_ptr).to(tl.bfloat16) |
| v_block = tl.load(v_block_ptr).to(tl.bfloat16) |
|
|
| grad_k_accum = tl.zeros([kv_block_size, feature_dim], dtype=tl.float32) |
| grad_v_accum = tl.zeros([kv_block_size, feature_dim], dtype=tl.float32) |
|
|
| for q_block_id in range(num_blocks_in_seq): |
| is_connected = tl.load(block_mask_ptr + fine_mask_start + q_block_id) |
|
|
| if is_connected: |
| for tile_in_block in range(tiles_per_block): |
| tile_idx = q_block_id * tiles_per_block + tile_in_block |
| q_tile_start = tile_idx * q_tile_size |
|
|
| q_tile_ptr = tl.make_block_ptr( |
| base=q_ptr + (batch_offset + q_tile_start) * num_q_heads * feature_dim, |
| shape=(q_tile_size, num_q_heads, feature_dim), |
| strides=(num_q_heads * feature_dim, feature_dim, 1), |
| offsets=(0, kv_head_idx * q_heads_per_kv_head, 0), |
| block_shape=(q_tile_size, q_heads_per_kv_head, feature_dim), |
| order=(0, 1, 2), |
| ) |
|
|
| grad_o_tile_ptr = tl.make_block_ptr( |
| base=grad_o_ptr + (batch_offset + q_tile_start) * num_q_heads * feature_dim, |
| shape=(q_tile_size, num_q_heads, feature_dim), |
| strides=(num_q_heads * feature_dim, feature_dim, 1), |
| offsets=(0, kv_head_idx * q_heads_per_kv_head, 0), |
| block_shape=(q_tile_size, q_heads_per_kv_head, feature_dim), |
| order=(0, 1, 2), |
| ) |
|
|
| lse_tile_ptr = tl.make_block_ptr( |
| base=lse_ptr + (batch_offset + q_tile_start) * num_q_heads, |
| shape=(q_tile_size, num_q_heads), |
| strides=(num_q_heads, 1), |
| offsets=(0, kv_head_idx * q_heads_per_kv_head), |
| block_shape=(q_tile_size, q_heads_per_kv_head), |
| order=(1, 0), |
| ) |
|
|
| delta_tile_ptr = tl.make_block_ptr( |
| base=delta_ptr + (batch_offset + q_tile_start) * num_q_heads, |
| shape=(q_tile_size, num_q_heads), |
| strides=(num_q_heads, 1), |
| offsets=(0, kv_head_idx * q_heads_per_kv_head), |
| block_shape=(q_tile_size, q_heads_per_kv_head), |
| order=(1, 0), |
| ) |
|
|
| q_tile = tl.load(q_tile_ptr) * softmax_scale * LOG2_E |
| q_tile = tl.reshape(q_tile, (q_tile_size * q_heads_per_kv_head, feature_dim)) |
| q_tile = q_tile.to(tl.bfloat16) |
|
|
| grad_o_block = tl.load(grad_o_tile_ptr) |
| grad_o_block = tl.reshape(grad_o_block, (q_tile_size * q_heads_per_kv_head, feature_dim)) |
| grad_o_block = grad_o_block.to(tl.bfloat16) |
|
|
| lse_vals = tl.load(lse_tile_ptr) |
| lse_vals = tl.reshape(lse_vals, (q_tile_size * q_heads_per_kv_head,)) |
|
|
| delta_vals = tl.load(delta_tile_ptr) |
| delta_vals = tl.reshape(delta_vals, (q_tile_size * q_heads_per_kv_head,)) |
|
|
| attention_scores = tl.dot(k_block, tl.trans(q_tile)) |
| attention_probs = tl.exp2(attention_scores - lse_vals[None, :]) |
| grad_v_accum += tl.dot(attention_probs.to(tl.bfloat16), grad_o_block) |
| grad_times_v = tl.dot(v_block, tl.trans(grad_o_block)) |
| grad_scores = attention_probs * (grad_times_v - delta_vals[None, :]) * LN_2 |
| grad_k_accum += tl.dot(grad_scores.to(tl.bfloat16), q_tile) |
|
|
| tl.store(grad_k_ptr, grad_k_accum.to(grad_k_ptr.dtype.element_ty)) |
| tl.store(grad_v_ptr, grad_v_accum.to(grad_v_ptr.dtype.element_ty)) |
|
|
|
|
| def mosaic_attn_bwd( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| output: torch.Tensor, |
| lse: torch.Tensor, |
| grad_o: torch.Tensor, |
| softmax_scale: float, |
| block_indices: torch.LongTensor, |
| block_size: int, |
| ): |
| batch_size, seq_len, num_kv_heads, feature_dim = k.shape |
| num_q_heads = q.shape[2] |
| num_kv_blocks_per_q_block = block_indices.shape[-1] |
| q_heads_per_kv_head = num_q_heads // num_kv_heads |
| num_blocks_in_seq = seq_len // block_size |
|
|
| grad_q = torch.empty_like(q) |
| grad_k = torch.empty_like(k) |
| grad_v = torch.empty_like(v) |
|
|
| block_mask = mosaic_block_mask(block_indices) |
|
|
| delta = (output * grad_o).sum(dim=-1) |
|
|
| grid_dq = lambda META: ( |
| triton.cdiv(seq_len, META['q_tile_size']), |
| q_heads_per_kv_head, |
| batch_size * num_kv_heads |
| ) |
|
|
| mosaic_attn_bwd_q_kernel[grid_dq]( |
| q_ptr=q, |
| k_ptr=k, |
| v_ptr=v, |
| lse_ptr=lse, |
| delta_ptr=delta, |
| grad_o_ptr=grad_o, |
| grad_q_ptr=grad_q, |
| block_indices_ptr=block_indices, |
| softmax_scale=softmax_scale, |
| seq_len=seq_len, |
| num_kv_heads=num_kv_heads, |
| num_q_heads=num_q_heads, |
| q_heads_per_kv_head=q_heads_per_kv_head, |
| feature_dim=feature_dim, |
| kv_block_size=block_size, |
| num_kv_blocks_per_q_block=num_kv_blocks_per_q_block, |
| ) |
|
|
| grid_dkv = (num_blocks_in_seq, batch_size * num_kv_heads) |
|
|
| mosaic_attn_bwd_kv_kernel[grid_dkv]( |
| q_ptr=q, |
| k_ptr=k, |
| v_ptr=v, |
| lse_ptr=lse, |
| delta_ptr=delta, |
| grad_o_ptr=grad_o, |
| grad_k_ptr=grad_k, |
| grad_v_ptr=grad_v, |
| block_mask_ptr=block_mask, |
| softmax_scale=softmax_scale, |
| seq_len=seq_len, |
| num_kv_heads=num_kv_heads, |
| num_q_heads=num_q_heads, |
| q_heads_per_kv_head=q_heads_per_kv_head, |
| feature_dim=feature_dim, |
| kv_block_size=block_size, |
| ) |
|
|
| return grad_q, grad_k, grad_v |
|
|
|
|
| class MosaicAttnFunction(torch.autograd.Function): |
|
|
| @staticmethod |
| @torch.amp.custom_fwd(device_type='cuda') |
| def forward( |
| ctx: torch.autograd.function.FunctionCtx, |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| block_indices: torch.Tensor, |
| block_size: int, |
| softmax_scale: float |
| ): |
| q, k, v, block_indices = map(lambda x: x.contiguous(), (q, k, v, block_indices)) |
|
|
| ctx.dtype = q.dtype |
|
|
| output, lse = mosaic_attn_fwd( |
| q=q, k=k, v=v, |
| block_indices=block_indices, |
| block_size=block_size, |
| softmax_scale=softmax_scale, |
| ) |
|
|
| ctx.save_for_backward(q, k, v, output, lse, block_indices) |
| ctx.block_size = block_size |
| ctx.softmax_scale = softmax_scale |
|
|
| return output.to(q.dtype) |
|
|
| @staticmethod |
| @torch.amp.custom_bwd(device_type='cuda') |
| def backward( |
| ctx: torch.autograd.function.FunctionCtx, |
| grad_o: torch.Tensor |
| ): |
| q, k, v, output, lse, block_indices = ctx.saved_tensors |
| grad_o = grad_o.contiguous() |
| grad_q, grad_k, grad_v = mosaic_attn_bwd( |
| q=q, k=k, v=v, output=output, lse=lse, grad_o=grad_o, |
| softmax_scale=ctx.softmax_scale, |
| block_indices=block_indices, |
| block_size=ctx.block_size, |
| ) |
| return grad_q, grad_k, grad_v, None, None, None |
|
|
|
|
| def mosaic_sparse_attn( |
| q: torch.Tensor, |
| k: torch.Tensor, |
| v: torch.Tensor, |
| block_indices: torch.LongTensor, |
| block_size: int, |
| softmax_scale: float = None, |
| ): |
| softmax_scale = q.shape[-1] ** -0.5 if softmax_scale is None else softmax_scale |
| return MosaicAttnFunction.apply(q, k, v, block_indices, block_size, softmax_scale) |
|
|
