Kasra Mazaheri

Scaling State Space Models with Block-Sparsity and Fused Kernels

A scaling story for oscillator state-space layers: block-sparse projection heads control dense coupling, while IO-aware FlashDOSS kernel fuses projection, scan, and projection-back work to avoid expensive state-domain traffic.

Flash MoBA: Optimizing Mixture of Block Attention

Guangxuan Xiao*, Junxian Guo*, Kasra Mazaheri, Song Han

Flash MoBA studies how mixture-of-block attention can route long-context computation sparsely while preserving quality, connecting routing accuracy, block size, local key aggregation, and efficient CUDA execution.

[paper] [code]

Broteina: Self-Distillation Unlocks Few-Step Protein Design

Kasra Mazaheri and collaborators

Preprint available to share

Broteina addresses train and test-time score misalignment in protein backbone diffusion models by first self-distilling the low-temperature sampling distribution, then distilling these 400-step models into single- and few-step generators.

FlashLinOSS: Expressive, Scalable, and Efficient Architectures for Oscillatory State-Space Models

MIT Distributed Robotics Lab

Preprint available to share

FlashLinOSS shows that block-sparse oscillator SSMs can outperform denser variants with fewer projection parameters, while IO-aware fused kernels reduce runtime by up to 7.8x and peak memory by about 3x.

[code]