# # Copyright (c) 2026 BEL ESPRIT D ACCORD TRUST HOLDINGS INC # All rights reserved. # SPDX-License-Identifier: Apache-2.0 # Copyright 2026 X.AI Corp. """ 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 # ============================================================================ # PRIMITIVE: Backward Softmax State # ============================================================================ @dataclass class BackwardState: dq_acc: jax.Array lse: jax.Array delta: jax.Array # ============================================================================ # MOSAIC CAP HELPERS (return aux for efficient gradient) # ============================================================================ 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}") # ============================================================================ # RANKER MASK (KV-Q orientation for dK/dV kernel) # ============================================================================ 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 # ============================================================================ # KERNEL: Mosaic dQ (Query Gradient) # ============================================================================ 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 # S = Q @ K^T 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) # P = exp2(S_capped * log2e - LSE) 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])) # Ranker mask 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) # dp = do @ V^T 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 - delta) ds = p * (dp - lax.broadcast_in_dim(delta, (block_q, block_kv), [0])) # Z-loss 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] ) # Cap gradient 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 # dQ += ds @ K 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 # ============================================================================ # KERNEL: Mosaic dK/dV (Key/Value Gradients) # ============================================================================ 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 # S^T = K @ Q^T 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) # P^T = exp2(S^T * log2e - LSE) 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) # dV += P^T @ do, dpT = V @ do^T 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) # dsT = pT * (dpT - delta) 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 # dK += dsT @ Q 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 # ============================================================================ # PREPROCESS BACKWARD # ============================================================================ 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 # ============================================================================ # BACKWARD PASS ORCHESTRATION # ============================================================================ @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 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) # dK/dV kernel 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) # GQA reduction # dk, dv shapes: (B, KV, H, D) -> reduce to (B, KV, H_kv, D) # ... dispatch kernels and reduce ... 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." ) # ============================================================================ # TEST HARNESS # ============================================================================ 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()