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