[Bounty $1500] Diffusion Drive bring up using TTNN APIs

Tenstorrent Bounties

Issue ID: I_kwDOI9Wqc87UGpYP

:memo: Background

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.

:bullseye: 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).
  • Store intermediates in L1 where beneficial.
  • Use recommended TTNN/tt-metal CNN flows.
  • Use TT library of fused ops like relu with conv

Stage 3 β€” Deeper Optimization

  • Maximize core counts used per inference.
  • Implement deeper TT-specific optimizations (e.g., grouped/depthwise convs, reordering to avoid unnecessary Tensor manipulations).
  • Minimize memory and TM/Tensor Manipulation overheads.
  • Document any advanced tuning or known issues.

:compass: Guidance & Starting Points

:magnifying_glass_tilted_right: 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.

:bar_chart: 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.

  • Functional model implementation
  • Validation logs (output correctness)
  • Performance report + header for final review

Links:


:books: Resources