This bounty is for bringing up DiffusionDrive
using TTNN APIs on Tenstorrent hardware (Wormhole or Blackhole).
DiffusionDrive is a truncated diffusion model designed for real-time end-to-end autonomous driving.
It replaces the slow, multi-step denoising process of conventional diffusion policies with a lightweight anchored and truncated diffusion process, achieving over 10Γ faster inference while maintaining multi-modal driving behavior and state-of-the-art accuracy.
Without any extra training tricks, DiffusionDrive reaches 88.1 PDMS on the NAVSIM benchmark using the same ResNet-34 backbone, trained via imitation learning on human driving trajectories, and runs at real-time 45 FPS.
For full architectural and algorithmic details, refer to the paper
.
The goal of this bounty is to enable efficient DiffusionDrive inference on Tenstorrent hardware, allowing users and partners to leverage high-throughput, low-latency execution for end-to-end driving policy and trajectory prediction.
What Success Looks Like
A successful submission will fulfill all requirements in the following stages. Payout is made after all three stages are completed.
Stage 1 β Bring-Up
Implement Diffusion Drive using TTNN APIs (Python).
Model runs on either N150 or N300 hardware with no errors.
Produces valid outputs/ego vehicle trajectories matching the implementation here.
Output is verifiable (visualization + simple metric checks).
Clear instructions for setup and running the model.
Stage 2 β Basic Optimizations
Use optimal sharded/interleaved memory configs and sharding strategy for conv ops in backbone/heads.
Fuse simple ops where possible (e.g., Conv+BN+Activation using TT fused ops).
Refer to TT fused ops for opportunities to optimize.
Target input resolutions: start with 640Γ640Γ3 (or model default).
Ask for help or file issues if ops are missing in TTNN.
Possible Approaches
Start from an existing implementation of the model /the link shared above and port layers one by one to TTNN.
Validate each submoduleβs output against CPU/PyTorch reference before full integration.
Experiment with different sharding strategies and memory configs for convolutions.
Use TTNN profiling tools to identify bottlenecks and areas for fusion.
Open a draft PR early to get feedback on your approach.
Result Submission Guidelines
Beyond the model implementation itself. Contributors must submit the following material as a proof of work. However, feel free to open a PR at any time if you want us checking you are on the right track. Just understand that payout is only made after all 3 stages are completed.