Add StreamIndex Kernel for CSA [Deepseek v4] - #5079

Draft

octatrifan wants to merge 13 commits into

mainfrom

octatrifan-dsv4-streamindex

Commits on Aug 24, 2026

Add fused Pallas TPU CSA StreamIndex score kernel for DeepSeek-V4
octatrifan
committed
f03759f
Optimize CSA StreamIndex kernel tile sizes and add multi-device shard_map support
octatrifan
committed
2e0752d
Optimize CSA StreamIndex kernel with VMEM head chunking and seq-major RoPE
octatrifan
committed
195e958

Commits on Aug 29, 2026

Support training backward pass via custom VJP in CSA StreamIndex
octatrifan
committed
45aeaee
Add csa_streamindex_score_head_major kernel and dispatch
octatrifan
committed
da439a9
Update unit tests for csa_streamindex_score_head_major
octatrifan
committed
fbe3672

Commits on Aug 31, 2026