File size: 1,585 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 | #
# Copyright (c) 2026 BEL ESPRIT D ACCORD TRUST HOLDINGS INC
# All rights reserved.
# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 X.AI Corp.
"""
Online Softmax State for Flash Attention.
Invariant: l = sum(exp(qk - m)), m = max(qk) per query.
"""
from dataclasses import dataclass
from typing import Optional
import jax
import jax.numpy as jnp
from .cap_functions import CapMethod, CapParams, cap_forward
@dataclass
class SoftmaxState:
m: jax.Array
l: jax.Array
acc: jax.Array
@staticmethod
def init(block_q: int, head_dim: int, dtype=jnp.float32) -> "SoftmaxState":
return SoftmaxState(
m=jnp.full((block_q,), -jnp.inf, dtype=dtype),
l=jnp.zeros((block_q,), dtype=dtype),
acc=jnp.zeros((block_q, head_dim), dtype=dtype)
)
def update(self, qk: jax.Array, v: jax.Array,
cap_method: CapMethod, cap_params: CapParams,
temp: Optional[jax.Array] = None) -> "SoftmaxState":
if temp is not None:
qk = qk * temp[..., None]
qk_capped = cap_forward(qk, cap_method, cap_params)
m_new = jnp.maximum(self.m, jnp.max(qk_capped, axis=1))
alpha = jnp.exp(self.m - m_new)
l_new = self.l * alpha + jnp.sum(jnp.exp(qk_capped - m_new[:, None]), axis=1)
p = jnp.exp(qk_capped - m_new[:, None])
p = p / l_new[:, None]
acc_new = self.acc * alpha[:, None] + p @ v
return SoftmaxState(m=m_new, l=l_new, acc=acc_new)
def finalize(self) -> jax.Array:
return self.acc / self.l[:, None]
|