Evo-1: avoid redundant vision QKV materialization

[j0yk1ll]

Title

Evo-1: avoid redundant vision QKV materialization

Summary

Adds an opt-in Evo-1 vision optimization that lets flash attention consume the strided Q/K/V layouts produced by reshape + permute directly, instead of materializing contiguous F32 copies in every InternViT layer.

Enable with:

VLA_EVO1_VISION_STRIDED_QKV=1

The optimization is gated on flash attention being enabled. The existing contiguous path remains the default, and the non-flash attention path is unchanged.

Motivation

Evo-1's InternViT has 24 transformer layers. In the existing flash-attention path, each layer materializes contiguous Q, K, and V tensors after reshape + permute.

That results in three redundant F32 materializations per layer, or 72 per camera view.

CUDA flash attention can consume these strided layouts directly, so the extra materializations are unnecessary.

Performance

Tested on Jetson Orin Nano Super, CUDA 12.6, with:

• Evo-1 LIBERO GGUF

• 3 camera views

• 448×448 input

• BF16 activations

• flash attention enabled

• CUDA graphs enabled

• horizon 50

• 32 flow steps

Matched steady-state A/B:

  | Baseline | Strided QKV | Delta -- | -- | -- | -- Total latency | 1267.6 ms | 1174.9 ms | -92.7 ms (-7.3%) Vision latency | 725.7 ms | 632.0 ms | -93.7 ms (-12.9%)

A separate 5-repetition optimized run measured approximately 1171.2 ms total and 630.5 ms vision.

The non-vision portion is effectively unchanged.

Numerical validation

A deterministic production-configuration A/B was run with 3×448 inputs, horizon 50, and all 32 flow steps.

The complete action output contains 1200 float32 values (50 × 24).

Baseline versus strided-QKV output:

elements      : 1200
max_abs_diff  : 0.0
mean_abs_diff : 0.0
RMSE          : 0.0
exact_equal   : True
byte_equal    : True

The optimized path therefore produced byte-identical final actions in this deterministic test.

Compatibility

The optimization is opt-in and only activates when flash attention is enabled:

VLA_EVO1_VISION_STRIDED_QKV=1

Without the environment variable, the existing contiguous Q/K/V path is retained.

When flash attention is disabled, the environment variable does not change the Q/K behavior and the existing non-flash V layout/materialization path remains in use.