Meta proposed a new attention kernel for generative recommender original paper.
HSTU Multi-Head Attention Definition
For each packed sequence and attention head $h$, let
$Q_h, K_h \in \mathbb{R}^{L \times d}$ and
$V_h \in \mathbb{R}^{L \times d_v}$. HSTU attention is defined as:
$$
S_h = \alpha Q_h K_h^{T} + R_h,
$$
$$
A_h = \frac{1}{L_{\mathrm{scale}}}
\left(M \odot \mathrm{SiLU}(S_h)\right),
$$
$$
O_h = A_h V_h,
$$
$$
\mathrm{HSTU\text{-}MHA}(Q,K,V)
= \mathrm{Concat}(O_1,O_2,\ldots,O_H).
$$
Here:
-
$\mathrm{SiLU}(x)=x,\sigma(x)=\frac{x}{1+e^{-x}}$.
-
$\alpha$ is the QK score scale.
-
$M_{ij}\in{0,1}$ is the attention mask: 1 for a valid QK pair and 0
otherwise. It can represent causal, local/sliding-window, target-group, or
arbitrary masks.
-
$R_h$ is an optional relative attention bias (RAB).
-
$L_{\mathrm{scale}}$ is the sequence normalization factor exposed as
scaling_seqlen; its default public behavior is max_seqlen_q.
Unlike standard Transformer attention, HSTU applies SiLU to the QK scores and
does not use row-wise softmax normalization.
We need to implement it with CuTeDSL on Blackwell.
Meta proposed a new attention kernel for generative recommender original paper.
HSTU Multi-Head Attention Definition
For each packed sequence and attention head$h$ , let
$Q_h, K_h \in \mathbb{R}^{L \times d}$ and
$V_h \in \mathbb{R}^{L \times d_v}$ . HSTU attention is defined as:
Here:
otherwise. It can represent causal, local/sliding-window, target-group, or
arbitrary masks.
scaling_seqlen; its default public behavior ismax_seqlen_q.Unlike standard Transformer attention, HSTU applies SiLU to the QK scores and
does not use row-wise softmax normalization.
We need to implement it with CuTeDSL on Blackwell.