perf(otfm): encode condition once, batch conditions in predict#309
Open
selmanozleyen wants to merge 1 commit into
Open
perf(otfm): encode condition once, batch conditions in predict#309selmanozleyen wants to merge 1 commit into
selmanozleyen wants to merge 1 commit into
Conversation
selmanozleyen
force-pushed
the
perf/predict-encode-once-batched
branch
from
July 14, 2026 13:51
3e27cfb to
f532232
Compare
Split ConditionalVelocityField.__call__ into encode_condition + velocity_from_embedding so the predict ODE integrates a precomputed embedding instead of re-encoding the condition on every step, and make OTFM.predict solve all conditions in one condition-batched vmap (no guidance) instead of a per-condition Python loop, concatenating ragged per-condition cells so no padding is needed. Numerically identical to the previous per-condition/per-step path (max abs diff ~1e-6, floating point only); the guidance path is unchanged. Encoding once is ~1.2x with attention pooling; batching conditions is ~1.5x and compiles one kernel instead of recompiling per distinct source-cell count.
selmanozleyen
force-pushed
the
perf/predict-encode-once-batched
branch
from
July 14, 2026 13:54
7fdef37 to
d7fa03a
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Split ConditionalVelocityField.call into encode_condition + velocity_from_embedding so the predict ODE integrates a precomputed embedding instead of re-encoding the condition on every step, and make OTFM.predict solve all conditions in one condition-batched vmap (no guidance) instead of a per-condition Python loop, concatenating ragged per-condition cells so no padding is needed.
Numerically identical to the previous per-condition/per-step path (max abs diff ~1e-6, floating point only); the guidance path is unchanged. Encoding once is ~1.2x with attention pooling; batching conditions is ~1.5x and compiles one kernel instead of recompiling per distinct source-cell count.