Skip to content

feat: add MojoPaddedWindowAttention and MojoConv1d.#274

Open
wwens7 wants to merge 10 commits into
masterfrom
wens/audio_core
Open

feat: add MojoPaddedWindowAttention and MojoConv1d.#274
wwens7 wants to merge 10 commits into
masterfrom
wens/audio_core

Conversation

@wwens7

@wwens7 wwens7 commented May 6, 2026

Copy link
Copy Markdown
Collaborator

No description provided.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces the MojoPaddedWindowAttention and MojoConv1d operators, including a high-performance Triton kernel implementation for padded window attention on NPU backends. The changes also include reference implementations, core operator registrations, and comprehensive accuracy and performance tests. The review feedback focuses on improving the Triton kernel's robustness and performance, specifically by addressing a potential division-by-zero risk, optimizing loop ranges, and ensuring power-of-2 block dimensions. Additionally, there are suggestions to reduce computational overhead in the operator's forward pass by removing redundant memory copies and optimizing the reference implementation's vectorization.

Comment thread mojo_opset/backends/ttx/kernels/npu/conformer_attention.py Outdated
Comment thread mojo_opset/backends/ttx/kernels/npu/padded_window_attention.py Outdated
Comment thread mojo_opset/backends/ttx/kernels/npu/conformer_attention.py Outdated
Comment thread mojo_opset/backends/ttx/operators/attention.py Outdated
Comment thread mojo_opset/backends/ttx/operators/attention.py Outdated
Comment thread mojo_opset/core/operators/attention.py Outdated
Comment thread mojo_opset/core/operators/attention.py Outdated
Comment thread mojo_opset/core/operators/attention.py Outdated
@wwens7 wwens7 force-pushed the wens/audio_core branch from 8e3f740 to 414cd20 Compare May 8, 2026 03:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants