[FA4/B200] return_lse=True 使 D64 向前运行速度提高了约 19%

作者: drisspg创建于 2026年7月25日更新于 2026年9月6日

def make_call( interface: ModuleType, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, return_lse: bool, ) -> Callable[[], torch.Tensor]: """Create one allocation-free FA4 forward specialization.""" out = torch.empty_like(q) lse = ( torch.empty(q.shape[0], q.shape[2], q.shape[1], device=q.device) if return_lse else None )

def run() -> torch.Tensor:
    return

内容来源: Dao-AILab/flash-attention