File size: 2,282 Bytes
677e207 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | #
# 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()
|