Commit d9c3d2f

authored and

recml authors

committed

Add pointwise splash attention TPU kernel for HSTU.

PiperOrigin-RevId: 761337655

1 parent bb3d157 commit d9c3d2fCopy full SHA for d9c3d2f

2 files changed

+1133

-0

lines changed

recml/core/ops
- hstu_ops.py
- hstu_ops_test.py

2 files changed

+1133

-0

lines changed

Comments

(0)