| |
| |
| |
|
|
| |
| |
| """ |
| Mosaic GPU Backward Kernels for Unified Attention with ALSC Cap. |
| Warp-specialized: 2 compute WGs + 1 memory WG per kernel. |
| |
| Backward derivatives (chain rule): |
| do -> dp = do @ V^T |
| ds = p * (dp - delta) where delta = sum(p * dp) = sum(o * do) / l |
| ds *= temp |
| ds *= cap_grad(qk) # ALSC: alpha/cap * sigma'(z) / denom |
| ds *= sm_scale |
| dQ = ds @ K |
| dK = ds^T @ Q |
| dV = p^T @ do |
| |
| Z-loss gradient: ds += z_weight * p * (log(l) + m) / l |
| """ |
|
|
| import functools |
| import math |
| from dataclasses import dataclass |
| from typing import Optional |
|
|
| import jax |
| import jax.numpy as jnp |
| from jax import lax |
| from jax.experimental import pallas as pl |
| from jax.experimental.pallas import mosaic_gpu as plgpu |
|
|
| from .cap_functions import CapMethod, CapParams, cap_forward, cap_grad, _sigmoid |
| from .segment_bounds import SegmentBounds |
| from .kernel_config import KernelConfig |
|
|
|
|
| |
| |
| |
|
|
| @dataclass |
| class BackwardState: |
| dq_acc: jax.Array |
| lse: jax.Array |
| delta: jax.Array |
|
|
|
|
| |
| |
| |
|
|
| def _apply_cap_mosaic(s, method: CapMethod, params: CapParams): |
| if method == "none" or params.cap <= 0: |
| return s, None, None |
| elif method == "tanh": |
| s_capped = params.cap * jnp.tanh(s / params.cap) |
| return s_capped, jnp.tanh(s / params.cap), None |
| elif method == "soft_sign": |
| abs_s = jnp.abs(s) |
| soft_sign = 1.0 / (1.0 + abs_s / params.cap) |
| s_capped = s * soft_sign |
| return s_capped, None, soft_sign |
| elif method == "alsc": |
| cap, alpha, beta = params.cap, params.alpha, params.beta |
| z = alpha * (s / cap + beta) |
| sigma_z = _sigmoid(z) |
| sigma_beta = _sigmoid(alpha * beta) |
| denom = sigma_beta * (1 - sigma_beta) + 1e-12 |
| s_capped = cap * (sigma_z - sigma_beta) / denom |
| return s_capped, sigma_z, None |
| else: |
| raise ValueError(f"Unknown cap method: {method}") |
|
|
|
|
| def _apply_cap_grad_mosaic(ds, s, method: CapMethod, params: CapParams, |
| qk_tanh, soft_sign): |
| if method == "none" or params.cap <= 0: |
| return ds |
| elif method == "tanh": |
| return ds * (1 - qk_tanh ** 2) |
| elif method == "soft_sign": |
| return ds * soft_sign ** 2 |
| elif method == "alsc": |
| cap, alpha, beta = params.cap, params.alpha, params.beta |
| sigma_z = qk_tanh |
| sigma_beta = _sigmoid(alpha * beta) |
| denom = sigma_beta * (1 - sigma_beta) + 1e-12 |
| grad = (alpha / cap) * sigma_z * (1 - sigma_z) / denom |
| return ds * grad |
| else: |
| raise ValueError(f"Unknown cap method: {method}") |
|
|
|
|
| |
| |
| |
|
|
| def ranker_mask_mosaic_kv_q(q_seq_base, block_q, kv_seq_base, block_kv, bounds): |
| kv_ids = plgpu.broadcasted_iota( |
| jnp.int32, (block_kv, block_q), 0, layout=plgpu.Layout.WGMMA) + kv_seq_base |
| q_ids = plgpu.broadcasted_iota( |
| jnp.int32, (block_kv, block_q), 1, layout=plgpu.Layout.WGMMA) + q_seq_base |
|
|
| q_hist = (q_ids >= bounds.history_lower) & (q_ids < bounds.history_upper) |
| q_cand = (q_ids >= bounds.candidate_lower) & (q_ids < bounds.candidate_upper) |
| kv_hist = (kv_ids >= bounds.history_lower) & (kv_ids < bounds.history_upper) |
| kv_cand = (kv_ids >= bounds.candidate_lower) & (kv_ids < bounds.candidate_upper) |
|
|
| hist_mask = kv_hist & (q_hist | q_cand) |
| cand_self = q_cand & kv_cand & (q_ids == kv_ids) |
| return hist_mask | cand_self |
|
|
|
|
| |
| |
| |
|
|
| def make_mosaic_backward_dq_kernel(config: KernelConfig, block_q: int, block_kv: int, |
| head_dim: int, q_heads_per_kv_head: int, |
| compute_wgs: int = 2, max_concurrent: int = 4): |
| dtype = jnp.bfloat16 |
|
|
| def kernel(q_ref, k_ref, v_ref, do_ref, lse_ref, delta_ref, bound_ref, dq_ref, |
| smem_buffers, buffer_barriers): |
| batch = lax.axis_index("batch") |
| q_head = lax.axis_index("heads") |
| wg_idx = lax.axis_index("wg") |
| q_seq = lax.axis_index("q_seq") |
|
|
| kv_head = lax.div(q_head, jnp.array(q_heads_per_kv_head, q_head.dtype)) |
|
|
| q_smem2, do_smem2, lse_smem2, delta_smem2 = smem_buffers |
| q_barriers, do_barriers, lse_barriers, delta_barriers = buffer_barriers |
|
|
| hl = plgpu.load(bound_ref, (batch, 0)) |
| hu = plgpu.load(bound_ref, (batch, 1)) |
| cl = plgpu.load(bound_ref, (batch, 2)) |
| cu = plgpu.load(bound_ref, (batch, 3)) |
| bounds = SegmentBounds(hl, hu, cl, cu) |
|
|
| q_tile_base = q_seq * (compute_wgs * block_q) |
| q_tile_end = q_tile_base + (compute_wgs * block_q) |
| q_seq_base = q_tile_base + wg_idx * block_q |
|
|
| hist_k_start = lax.div(hl, block_kv) |
| hist_k_end = pl.cdiv(hu, block_kv) |
| hist_steps = jnp.maximum(hist_k_end - hist_k_start, 0) |
|
|
| cand_start = jnp.maximum(cl, q_tile_base) |
| cand_end = jnp.minimum(cu, q_tile_end) |
| cand_has = cand_start < cand_end |
| cand_k_start = lax.div(cand_start, block_kv) |
| cand_k_end = pl.cdiv(cand_end, block_kv) |
| cand_steps = jnp.where(cand_has, cand_k_end - cand_k_start, 0) |
|
|
| total_kv_steps = hist_steps + cand_steps |
|
|
| def _compute_thread(pipeline_callback): |
| q_smem = q_smem2.at[wg_idx] |
| do_smem = do_smem2.at[wg_idx] |
| lse_smem = lse_smem2.at[wg_idx] |
| delta_smem = delta_smem2.at[wg_idx] |
|
|
| q_slice = (batch, pl.ds(q_seq_base, block_q), q_head) |
| plgpu.copy_gmem_to_smem(q_ref.at[q_slice], q_smem, q_barriers.at[wg_idx]) |
| plgpu.copy_gmem_to_smem(do_ref.at[q_slice], do_smem, do_barriers.at[wg_idx]) |
| plgpu.copy_gmem_to_smem( |
| delta_ref.at[batch, q_head, pl.ds(q_seq_base, block_q)], |
| delta_smem, delta_barriers.at[wg_idx] |
| ) |
| plgpu.copy_gmem_to_smem( |
| lse_ref.at[batch, q_head, pl.ds(q_seq_base, block_q)], |
| lse_smem, lse_barriers.at[wg_idx] |
| ) |
| for b in buffer_barriers: |
| plgpu.barrier_wait(b.at[wg_idx]) |
|
|
| delta = plgpu.load(delta_smem, (), layout=plgpu.Layout.WGMMA_ROW) |
| lse = plgpu.load(lse_smem, (), layout=plgpu.Layout.WGMMA_ROW) |
|
|
| dq_acc = plgpu.layout_cast( |
| jnp.full((block_q, head_dim), 0, dtype=jnp.float32), |
| plgpu.Layout.WGMMA |
| ) |
|
|
| final_dq = pipeline_callback((dq_acc, lse, delta)) |
|
|
| q_smem[...] = final_dq.astype(dtype) |
| plgpu.commit_smem() |
| plgpu.copy_smem_to_gmem(q_smem, dq_ref.at[q_slice]) |
| plgpu.wait_smem_to_gmem(0) |
|
|
| def kv_pipeline(kv_step, k_smem, v_smem, k_consumed, v_consumed, carry): |
| kv_step = kv_step[0] |
| q_smem, do_smem = q_smem2.at[wg_idx], do_smem2.at[wg_idx] |
| dq_acc, lse, delta = carry |
|
|
| kv_block_idx = jnp.where( |
| kv_step < hist_steps, |
| hist_k_start + kv_step, |
| cand_k_start + (kv_step - hist_steps) |
| ) |
| kv_seq_base = kv_block_idx * block_kv |
|
|
| |
| def compute_s(acc_ref): |
| plgpu.wgmma(acc_ref, q_smem, plgpu.transpose_ref(k_smem, (1, 0))) |
| return acc_ref[...] |
| s = pl.run_scoped(compute_s, plgpu.ACC((block_q, block_kv), jnp.float32)) |
|
|
| if config.sm_scale != 1.0: |
| s *= config.sm_scale |
| s_capped, qk_tanh, soft_sign = _apply_cap_mosaic( |
| s, config.cap_method, config.cap_params) |
|
|
| |
| log2e = math.log2(math.e) |
| s_scaled = s_capped * log2e |
| p = jnp.exp2(s_scaled - lax.broadcast_in_dim(lse, (block_q, block_kv), [0])) |
|
|
| |
| from .mosaic_kernel import ranker_mask_mosaic |
| mask = ranker_mask_mosaic(q_seq_base, block_q, kv_seq_base, block_kv, bounds) |
| p = jnp.where(mask, p, 0) |
|
|
| |
| def compute_dp(acc_ref): |
| plgpu.wgmma(acc_ref, do_smem, plgpu.transpose_ref(v_smem, (1, 0))) |
| return acc_ref[...] |
| dp = pl.run_scoped(compute_dp, plgpu.ACC((block_q, block_kv), jnp.float32)) |
| plgpu.barrier_arrive(v_consumed) |
|
|
| |
| ds = p * (dp - lax.broadcast_in_dim(delta, (block_q, block_kv), [0])) |
|
|
| |
| if config.z_loss_weight > 0.0: |
| logsumexp = lse * math.log(2.0) |
| ds += p * lax.broadcast_in_dim( |
| config.z_loss_weight * logsumexp, (block_q, block_kv), [0] |
| ) |
|
|
| |
| ds = _apply_cap_grad_mosaic( |
| ds, s, config.cap_method, config.cap_params, qk_tanh, soft_sign) |
|
|
| if config.sm_scale != 1.0: |
| ds *= config.sm_scale |
|
|
| |
| def compute_dq(acc_ref): |
| plgpu.wgmma(acc_ref, ds.astype(dtype), k_smem) |
| dq_acc = pl.run_state(compute_dq)(plgpu.ACC.init(dq_acc)) |
| plgpu.barrier_arrive(k_consumed) |
|
|
| return (dq_acc, lse, delta) |
|
|
| pipeline = plgpu.emit_pipeline_warp_specialized( |
| kv_pipeline, |
| grid=(total_kv_steps,), |
| max_concurrent_steps=max_concurrent, |
| num_compute_wgs=compute_wgs, |
| memory_registers=40, |
| wg_axis="wg", |
| manual_consumed_barriers=True, |
| compute_context=_compute_thread, |
| in_specs=[ |
| plgpu.BlockSpec( |
| block_shape=(block_kv, head_dim), |
| index_map=lambda i: ( |
| jnp.where(i < hist_steps, hist_k_start + i, |
| cand_k_start + (i - hist_steps)), 0), |
| transforms=[plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128)], |
| ), |
| plgpu.BlockSpec( |
| block_shape=(block_kv, head_dim), |
| index_map=lambda i: ( |
| jnp.where(i < hist_steps, hist_k_start + i, |
| cand_k_start + (i - hist_steps)), 0), |
| transforms=[plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128)], |
| ), |
| ], |
| ) |
|
|
| k_ref_slice = k_ref.at[batch, :, kv_head, :] |
| v_ref_slice = v_ref.at[batch, :, kv_head, :] |
| pipeline(k_ref_slice, v_ref_slice) |
|
|
| return kernel |
|
|
|
|
| |
| |
| |
|
|
| def make_mosaic_backward_dkv_kernel(config: KernelConfig, block_q: int, block_kv: int, |
| head_dim: int, q_heads_per_kv_head: int, |
| compute_wgs: int = 2, max_concurrent: int = 4): |
| dtype = jnp.bfloat16 |
|
|
| def kernel(q_ref, k_ref, v_ref, do_ref, lse_ref, delta_ref, bound_ref, |
| dk_ref, dv_ref, smem_buffers, buffer_barriers): |
| batch = lax.axis_index("batch") |
| q_head = lax.axis_index("heads") |
| wg_idx = lax.axis_index("wg") |
| kv_seq = lax.axis_index("kv_seq") |
|
|
| k_smem2, v_smem2 = smem_buffers |
| k_barriers, v_barriers = buffer_barriers |
|
|
| hl = plgpu.load(bound_ref, (batch, 0)) |
| hu = plgpu.load(bound_ref, (batch, 1)) |
| cl = plgpu.load(bound_ref, (batch, 2)) |
| cu = plgpu.load(bound_ref, (batch, 3)) |
| bounds = SegmentBounds(hl, hu, cl, cu) |
|
|
| kv_group_base = kv_seq * (compute_wgs * block_kv) |
| kv_group_end = kv_group_base + (compute_wgs * block_kv) |
| kv_seq_base = kv_group_base + wg_idx * block_kv |
|
|
| kv_is_hist = (kv_group_base < hu) & (kv_group_end > hl) |
|
|
| hist_q_start = lax.div(hl, block_q) |
| hist_q_end = pl.cdiv(hu, block_q) |
| hist_q_steps = jnp.maximum(hist_q_end - hist_q_start, 0) |
|
|
| cand_has = cl < cu |
| cand_q_start = lax.div(cl, block_q) |
| cand_q_end = pl.cdiv(cu, block_q) |
| cand_q_steps = jnp.where(cand_has, cand_q_end - cand_q_start, 0) |
|
|
| cand_overlap_start = jnp.maximum(cl, kv_group_base) |
| cand_overlap_end = jnp.minimum(cu, kv_group_end) |
| cand_overlap_has = cand_overlap_start < cand_overlap_end |
| cand_q_overlap_start = lax.div(cand_overlap_start, block_q) |
| cand_q_overlap_end = pl.cdiv(cand_overlap_end, block_q) |
| cand_q_overlap_steps = jnp.where( |
| cand_overlap_has, cand_q_overlap_end - cand_q_overlap_start, 0) |
|
|
| total_q_steps = jnp.where( |
| kv_is_hist, |
| hist_q_steps + cand_q_steps, |
| cand_q_overlap_steps |
| ) |
|
|
| def q_block_for_step(q_step): |
| q_step = jnp.asarray(q_step, jnp.int32) |
| hist_cand = jnp.where( |
| q_step < hist_q_steps, |
| hist_q_start + q_step, |
| cand_q_start + (q_step - hist_q_steps) |
| ) |
| cand_only = cand_q_overlap_start + q_step |
| return jnp.where(kv_is_hist, hist_cand, cand_only) |
|
|
| def _compute_thread(pipeline_callback): |
| k_smem = k_smem2.at[wg_idx] |
| v_smem = v_smem2.at[wg_idx] |
|
|
| kv_head = lax.div(q_head, jnp.array(q_heads_per_kv_head, q_head.dtype)) |
| plgpu.copy_gmem_to_smem( |
| k_ref.at[(batch, pl.ds(kv_seq_base, block_kv), kv_head)], |
| k_smem, k_barriers.at[wg_idx] |
| ) |
| plgpu.copy_gmem_to_smem( |
| v_ref.at[(batch, pl.ds(kv_seq_base, block_kv), kv_head)], |
| v_smem, v_barriers.at[wg_idx] |
| ) |
| plgpu.barrier_wait(k_barriers.at[wg_idx]) |
| plgpu.barrier_wait(v_barriers.at[wg_idx]) |
|
|
| dk_acc = plgpu.layout_cast( |
| jnp.full((block_kv, head_dim), 0, dtype=jnp.float32), plgpu.Layout.WGMMA) |
| dv_acc = plgpu.layout_cast( |
| jnp.full((block_kv, head_dim), 0, dtype=jnp.float32), plgpu.Layout.WGMMA) |
|
|
| dk, dv = pipeline_callback((dv_acc, dk_acc)) |
|
|
| k_smem[...] = dk.astype(dtype) |
| v_smem[...] = dv.astype(dtype) |
| plgpu.commit_smem() |
| plgpu.copy_smem_to_gmem( |
| k_smem, dk_ref.at[(batch, pl.ds(kv_seq_base, block_kv), q_head)], |
| commit_group=False) |
| plgpu.copy_smem_to_gmem( |
| v_smem, dv_ref.at[(batch, pl.ds(kv_seq_base, block_kv), q_head)], |
| commit_group=False) |
| plgpu.commit_smem_to_gmem_group() |
| plgpu.wait_smem_to_gmem(0) |
|
|
| def q_pipeline(q_step, q_smem, do_smem, lse_smem, delta_smem, |
| q_consumed, do_consumed, lse_consumed, delta_consumed, carry): |
| q_step = q_step[0] |
| k_smem, v_smem = k_smem2.at[wg_idx], v_smem2.at[wg_idx] |
| dk_acc, dv_acc = carry |
|
|
| q_block_idx = q_block_for_step(q_step) |
| q_seq_base_local = q_block_idx * block_q |
|
|
| |
| def compute_st(acc_ref): |
| plgpu.wgmma(acc_ref, k_smem, plgpu.transpose_ref(q_smem, (1, 0))) |
| return acc_ref[...] |
| sT = pl.run_scoped(compute_st, plgpu.ACC((block_kv, block_q), jnp.float32)) |
|
|
| if config.sm_scale != 1.0: |
| sT *= config.sm_scale |
| sT_capped, qk_tanh, soft_sign = _apply_cap_mosaic( |
| sT, config.cap_method, config.cap_params) |
|
|
| |
| log2e = math.log2(math.e) |
| lse = plgpu.load(lse_smem, (), layout=plgpu.Layout.WGMMA_COL) |
| plgpu.barrier_arrive(lse_consumed) |
|
|
| pT = jnp.exp2( |
| sT_capped * log2e - |
| lax.broadcast_in_dim(lse, (block_kv, block_q), [1])) |
|
|
| mask = ranker_mask_mosaic_kv_q( |
| q_seq_base_local, block_q, kv_seq_base, block_kv, bounds) |
| pT = jnp.where(mask, pT, 0) |
|
|
| |
| def compute_dv_dpt(refs): |
| dv_acc_ref, dpT_acc_ref = refs |
| plgpu.wgmma(dv_acc_ref, pT.astype(dtype), do_smem) |
| plgpu.wgmma(dpT_acc_ref, v_smem, plgpu.transpose_ref(do_smem, (1, 0))) |
|
|
| zeros = plgpu.layout_cast( |
| jnp.full((block_kv, block_q), 0, jnp.float32), plgpu.Layout.WGMMA) |
| dv_acc, dpT = pl.run_state(compute_dv_dpt)( |
| (plgpu.ACC.init(dv_acc), plgpu.ACC.init(zeros))) |
| plgpu.barrier_arrive(do_consumed) |
|
|
| |
| delta = plgpu.load(delta_smem, (), layout=plgpu.Layout.WGMMA_COL) |
| plgpu.barrier_arrive(delta_consumed) |
|
|
| dsT = pT * (dpT - lax.broadcast_in_dim(delta, (block_kv, block_q), [1])) |
|
|
| if config.z_loss_weight > 0.0: |
| logsumexp = lse * math.log(2.0) |
| dsT += pT * lax.broadcast_in_dim( |
| config.z_loss_weight * logsumexp, (block_kv, block_q), [1]) |
|
|
| dsT = _apply_cap_grad_mosaic( |
| dsT, sT, config.cap_method, config.cap_params, qk_tanh, soft_sign) |
|
|
| if config.sm_scale != 1.0: |
| dsT *= config.sm_scale |
|
|
| |
| def compute_dk(acc_ref): |
| plgpu.wgmma(acc_ref, dsT.astype(dtype), q_smem) |
| dk_acc = pl.run_state(compute_dk)(plgpu.ACC.init(dk_acc)) |
| plgpu.barrier_arrive(q_consumed) |
|
|
| return (dk_acc, dv_acc) |
|
|
| pipeline = plgpu.emit_pipeline_warp_specialized( |
| q_pipeline, |
| grid=(total_q_steps,), |
| max_concurrent_steps=max_concurrent, |
| num_compute_wgs=compute_wgs, |
| memory_registers=40, |
| wg_axis="wg", |
| manual_consumed_barriers=True, |
| compute_context=_compute_thread, |
| in_specs=[ |
| plgpu.BlockSpec( |
| block_shape=(block_q, head_dim), |
| index_map=lambda i: (q_block_for_step(i), 0), |
| transforms=[plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128)], |
| ), |
| plgpu.BlockSpec( |
| block_shape=(block_q, head_dim), |
| index_map=lambda i: (q_block_for_step(i), 0), |
| transforms=[plgpu.TilingTransform((8, 64)), plgpu.SwizzleTransform(128)], |
| ), |
| plgpu.BlockSpec( |
| block_shape=(block_q,), |
| index_map=lambda i: (q_block_for_step(i),), |
| ), |
| plgpu.BlockSpec( |
| block_shape=(block_q,), |
| index_map=lambda i: (q_block_for_step(i),), |
| ), |
| ], |
| ) |
|
|
| q_ref_slice = q_ref.at[batch, :, q_head, :] |
| do_ref_slice = do_ref.at[batch, :, q_head, :] |
| lse_ref_slice = lse_ref.at[batch, q_head, :] |
| delta_ref_slice = delta_ref.at[batch, q_head, :] |
| pipeline(q_ref_slice, do_ref_slice, lse_ref_slice, delta_ref_slice) |
|
|
| return kernel |
|
|
|
|
| |
| |
| |
|
|
| def make_preprocess_backward_kernel(block_q: int): |
| def kernel(out_ref, do_ref, l_ref, do_scaled_ref, delta_ref): |
| pid = pl.program_id(0) |
| off = pl.ds(pid * block_q, block_q) |
|
|
| o = plgpu.load(out_ref, (off, slice(None))).astype(jnp.float32) |
| do = plgpu.load(do_ref, (off, slice(None))).astype(jnp.float32) |
| l = plgpu.load(l_ref, (off,)).astype(jnp.float32) |
|
|
| do_scaled = jnp.where(l[:, None] == 0, 0.0, do / l[:, None]) |
| delta = jnp.sum(o * do_scaled, axis=1) |
|
|
| plgpu.store(do_scaled_ref, (off, slice(None)), do_scaled.astype(do_scaled_ref.dtype)) |
| plgpu.store(delta_ref, (off,), delta.astype(delta_ref.dtype)) |
|
|
| return kernel |
|
|
|
|
| |
| |
| |
|
|
| @dataclass(frozen=True) |
| class BackwardConfig: |
| block_q_dq: int = 64 |
| block_kv_dq: int = 64 |
| block_q_dkv: int = 64 |
| block_kv_dkv: int = 64 |
| compute_wgs: int = 2 |
| max_concurrent: int = 4 |
|
|
|
|
| def mosaic_attention_backward( |
| q, k, v, temp, segment_ids, out, do, lse, config: KernelConfig, |
| bwd_config: BackwardConfig, |
| ): |
| """ |
| Full backward pass using Mosaic GPU kernels. |
| Returns: dq, dk, dv, dtemp, dsegment_ids |
| """ |
| B, Q, H, D = q.shape |
| _, KV, H_kv, _ = k.shape |
| q_h_per_kv = H // H_kv |
| dtype = q.dtype |
|
|
| |
| dq_kernel = make_mosaic_backward_dq_kernel( |
| config, bwd_config.block_q_dq, bwd_config.block_kv_dq, |
| D, q_h_per_kv, bwd_config.compute_wgs, bwd_config.max_concurrent) |
|
|
| |
| dkv_kernel = make_mosaic_backward_dkv_kernel( |
| config, bwd_config.block_q_dkv, bwd_config.block_kv_dkv, |
| D, q_h_per_kv, bwd_config.compute_wgs, bwd_config.max_concurrent) |
|
|
| |
| |
| |
|
|
| dtemp = jnp.zeros_like(temp) |
| dsegment = jnp.zeros_like(segment_ids, dtype=dtype) |
|
|
| raise NotImplementedError( |
| "Full orchestration requires Mosaic forward to return (l, m). " |
| "Use mosaic_attention_backward after implementing save_residuals in forward." |
| ) |
|
|
|
|
| |
| |
| |
|
|
| def test_mosaic_backward_correctness(): |
| """Numerical gradient check against JAX autodiff.""" |
| B, Q, KV, H, H_kv, D = 1, 128, 128, 4, 2, 64 |
| key = jax.random.PRNGKey(42) |
|
|
| q = jax.random.normal(key, (B, Q, H, D), dtype=jnp.bfloat16) |
| k = jax.random.normal(key, (B, KV, H_kv, D), dtype=jnp.bfloat16) |
| v = jax.random.normal(key, (B, KV, H_kv, D), dtype=jnp.bfloat16) |
| temp = jax.random.uniform(key, (B, Q), minval=0.8, maxval=1.2).astype(jnp.bfloat16) |
|
|
| from .segment_bounds import HISTORY_SEGMENT_ID, CANDIDATE_SEGMENT_ID |
| segment_ids = jnp.zeros((B, Q), dtype=jnp.int32) |
| segment_ids = segment_ids.at[:, :Q // 2].set(HISTORY_SEGMENT_ID) |
| segment_ids = segment_ids.at[:, Q // 2:].set(CANDIDATE_SEGMENT_ID) |
|
|
| config = KernelConfig( |
| cap_method="alsc", |
| cap_params=CapParams(cap=30.0, alpha=4.0, beta=0.5), |
| sm_scale=1.0 / math.sqrt(D), |
| backend="mosaic", |
| ) |
|
|
| from .attention import unified_attention_reference |
|
|
| def loss_fn(q, k, v): |
| out = unified_attention_reference(q, k, v, temp, segment_ids, config) |
| return jnp.sum(out ** 2) |
|
|
| dq_ref, dk_ref, dv_ref = jax.grad(loss_fn, argnums=(0, 1, 2))(q, k, v) |
|
|
| print("Test structure ready. Requires Mosaic forward implementation.") |
| print(f"Shapes: dq={dq_ref.shape}, dk={dk_ref.shape}, dv={dv_ref.shape}") |
|
|
|
|
| if __name__ == "__main__": |
| test_mosaic_backward_correctness() |
|
|
