python/xrex_unified/tests.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.
"""Tests and verification for unified attention kernel."""

import jax
import jax.numpy as jnp

from .cap_functions import (
    CapParams, _alsc_forward, _alsc_inverse, _alsc_grad, cap_forward
)
from .segment_bounds import SegmentBounds, HISTORY_SEGMENT_ID, CANDIDATE_SEGMENT_ID
from .kernel_config import KernelConfig
from .attention import unified_attention_reference


def test_cap_functions():
    """Verify cap function properties."""
    x = jnp.linspace(-100, 100, 1000)

    params = CapParams(cap=30.0, alpha=4.0, beta=0.5)
    y = _alsc_forward(x, params)

    assert jnp.all(y >= -1e-5) and jnp.all(y <= params.cap + 1e-5), "ALSC bound violation"
    assert jnp.all(jnp.diff(y) >= -1e-5), "ALSC not monotonic"

    x_recon = _alsc_inverse(y, params)
    assert jnp.allclose(x, x_recon, atol=1e-4), "ALSC inverse failed"

    grad_analytic = _alsc_grad(x, params)
    grad_auto = jax.grad(lambda x: _alsc_forward(x, params).sum())(x)
    assert jnp.allclose(grad_analytic, grad_auto, rtol=1e-4), "ALSC grad mismatch"

    print("All cap function tests passed")


def test_attention_equivalence():
    """Test kernel matches reference."""
    B, Q, KV, H, H_kv, D = 2, 256, 256, 8, 2, 128
    key = jax.random.PRNGKey(0)

    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.5, maxval=1.5).astype(jnp.bfloat16)

    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),
        backend="triton",
    )

    out_ref = unified_attention_reference(q, k, v, temp, segment_ids, config)

    print("Reference output shape:", out_ref.shape)
    print("Equivalence test structure ready")


if __name__ == "__main__":
    test_cap_functions()
    test_attention_equivalence()