[Feature] Reflect vortex prefill improvements from PR 77

[jstjohn]

Problem & Motivation

A user has recently contributed improvements to vortex's kernels, we have independently updated the kernels, so we should explore the current state of vortex and see what we can pull in.

BioNeMo Framework Version

main

Category

Inference

Proposed Solution

See updates in Zymrael/vortex#77 and implement any relevant improvements from there on top of our dynamic inference engine implementation.

Expected Benefits

Potentially a further 1.5x improvement in inference perf.

Code Example