[ENH] tft_v2: enable efficient attention by setting need_weights=False

[ramanbansal1]

Reference Issues/PRs

N/A

What does this implement/fix? Explain your changes.

This PR optimizes the self-attention computation in TFT_v2 by explicitly passing need_weights=False to nn.MultiheadAttention.

Previously, TFT_v2 computed the attention output but immediately discarded the attention weights (attended_output, _ = self.self_attention(...)). However, because need_weights defaults to True in PyTorch, this silently forced PyTorch to use its slow, memory-intensive math backend to explicitly construct the $(L+H) \times (L+H)$ attention probability matrix.

By explicitly setting need_weights=False, we allow PyTorch 2.0+ to automatically route the computation through scaled_dot_product_attention (SDPA), unlocking FlashAttention and Memory-Efficient Attention kernels. This is a zero-behavior-change optimization that yields massive memory savings, especially for long sequence lengths (e.g. large max_encoder_length), which is typically the memory bottleneck for TFT.

What should a reviewer concentrate their feedback on?

• This is a zero-behavior-change optimization. The only change is how PyTorch allocates memory under the hood since the original attention weights were discarded anyway.

Did you add any tests for the change?

• No, but I benchmarked the layer to confirm the memory and speed improvements.

Benchmark Results (Tesla T4, Sequence Length = 4000, Batch Size = 32):

• Before (need_weights=True):
  • Peak Memory: 1106.06 MB
  • Avg Forward Time: 10.44 ms
• After (need_weights=False):
  • Peak Memory: 103.72 MB (~10x reduction)
  • Avg Forward Time: 5.62 ms (~2x faster)

Any other comments?

This brings the TFT_v2 attention logic in line with PyTorch best practices for memory-efficient attention.

PR checklist

• The PR title starts with either [ENH], [MNT], [DOC], or [BUG]. [ENH]
• Added/modified tests (Not applicable for this minor performance flag)
• Used pre-commit hooks when committing to ensure that code is compliant with hooks. Install hooks with pre-commit install.
  To run hooks independent of commit, execute pre-commit run --all-files

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (main@ec90f3e). Learn more about missing BASE report.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2429   +/-   ##
=======================================
  Coverage        ?   88.54%           
=======================================
  Files           ?      208           
  Lines           ?    11381           
  Branches        ?        0           
=======================================
  Hits            ?    10077           
  Misses          ?     1304           
  Partials        ?        0           

Flag	Coverage Δ
cpu	88.54% <100.00%> (?)
pytest	88.54% <100.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[phoeenniixx]

Question: Would it be better to not "hard code" the param, instead let the user do it?