python/xrex_unified/mosaic_backward.py

#
# 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()