Hyena Inference Updates to support Flash Decode

[jstjohn]

Description

• Support for Flash Decode at inference time in evo2. The changes had been made upstream in NeMo previously ( Hyena support for flash decode API NVIDIA-NeMo/NeMo#14315).
• Do not yet support cudagraph, which is another potential performance gain at inference time.
• Temporary work around for the issue of increased memory usage whenever passing an inference_context to hyena introduced by the flash decode API change: [BUG] NeMo upgrade supporting hyena flash decode requires more memory for generation #1013. This unblocks bumping nemo related PRs such as Bump NeMo #1012.
• All of this is related to supporting a future NIM backed by bionemo. Unfortunately this is demonstrating how much code overhead there is in adding this kind of support, and with the memory usage issues, the speed improvement may not be worth it.
• Training performance, as well as forward pass without inference context performance is not impacted.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 90.90909% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 83.34%. Comparing base (c03dc80) to head (103d76f).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
...ackages/bionemo-evo2/src/bionemo/evo2/run/infer.py	90.90%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1000      +/-   ##
==========================================
+ Coverage   83.31%   83.34%   +0.02%     
==========================================
  Files         148      148              
  Lines        9758     9766       +8     
==========================================
+ Hits         8130     8139       +9     
+ Misses       1628     1627       -1     

Files with missing lines	Coverage Δ
...ackages/bionemo-evo2/src/bionemo/evo2/run/infer.py	54.90% <90.90%> (+8.39%)	⬆️

... and 1 file with indirect coverage changes

[jstjohn]

TODO wait for NVIDIA-NeMo/NeMo#14315 to merge and upstream to Main NeMo which also has the change from NVIDIA-NeMo/NeMo#14359

Update: This is now done.

[skothenhill-nv]

I didn't see anything glaring. What would be helpful is more in the PR description. I'm missing a lot of context, why did we make these changes? how should we think about this if someone encounters a non-related bug in the future from this PR?

[farhadrgh]

LGTM