Skip to content

perf(vln): optimize ActiveVLN and StreamVLN inference - #12

Merged
ZhouAo-ZA merged 20 commits into
mainfrom
bench/activevln-rxr-oom-fix
Oct 28, 2026
Merged

ZhouAo-ZA merged 20 commits into
mainfrom
bench/activevln-rxr-oom-fix

Conversation

@yufoo1

@yufoo1 yufoo1 commented Oct 21, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Bring ActiveVLN and StreamVLN inference optimizations onto current main: speculative/action-tree decoding, GPU image/history processing, CUDA graphs, resident/incremental KV state, packed prefill, and shared fused operators. StreamVLN also gains bounded-memory INT8 execution and tensor-batch HTTP benchmark services.

Publish portable measured configurations and verified summaries for the same first 128 episodes of each fixed 256-episode selection under benchmarks/{activevln,streamvln}-benchmark/{rtx-4090,orin-nx}. RTX 4090 runs inference and rendering on one card; Orin NX uses local weights and the LAN 4090 for rendering. The NX ActiveVLN INT8 source completed only 164/256; its published subset is explicitly 128 episodes. NX StreamVLN retains the older v9 result. These are measurements of the recorded frozen sources, not a new GPU validation of this integrated tree.

The final diff contains model code, benchmark services, tests, profiles and documentation. Go2 deployment/navigation and asynchronous execution-prefix support remain on their separate branches. Raw logs, report archives and weights are excluded from the final tree. Habitat orchestration and SR/SPL calculation live in the companion EmbodiRun change.

Type of change

  • Kernel, layer, or numerical change
  • Documentation
  • Engine, scheduler, session, or rollout change (policy-local batching/serving only)

Verification

The selected ActiveVLN test changes remain excluded: 10 added files are removed and 3 existing files retain their main-branch versions. Residual grammar/batching cases that imported the removed test helper are also removed; independent grammar helpers and upload-layout cases remain. Inference code, benchmark implementations, profiles and results are unchanged.

Validated with the same CPU environment as CI: Python 3.12, Torch 2.10.0+cpu and Transformers 5.3.0.

CUDA_VISIBLE_DEVICES='' PYTHONPATH=. NO_PROXY=localhost,127.0.0.1 no_proxy=localhost,127.0.0.1 python -m pytest tests -q
ruff check embodiinfer tests
ruff format --check embodiinfer tests
python3 scripts/check_commit_policy.py --base upstream/main --head HEAD --title 'perf(vln): optimize ActiveVLN and StreamVLN inference'
git diff --check upstream/main HEAD

1141 passed, 311 skipped; full lint/format, public entrypoint and commit-policy checks passed. No new GPU, model-checkpoint or navigation workloads were started. Device tables retain the previously recorded frozen-source measurements.

Engine invariants

  • The engine remains model-neutral; model behavior stays in policies/models and reusable kernels in backends.
  • No deployment-runtime imports were added to the inference package.
  • Optional paths preserve explicit dtype, shape, state and numerical contracts; INT8 is documented as a separate precision experiment.
  • Dependency declarations and the lockfile are unchanged.

Checklist

  • Conventional Commit subjects and existing authors' DCO sign-offs are preserved; new commits use the configured human author and sign-off.
  • Proposals and benchmark/model documentation describe the implementation and validation limits.
  • No weights, datasets, raw benchmark outputs, recordings, credentials, private host addresses or private paths are added to the final diff.
  • Contribution follows the repository's Apache-2.0 license.

@yufoo1
yufoo1 requested a review from ZhouAo-ZA October 21, 2026 09:13
@yufoo1 yufoo1 changed the title Bench/activevln rxr oom fix perf(activevln): add batched inference optimizations Oct 21, 2026
yufoo1 and others added 16 commits October 21, 2026 18:03
Add policy-local CUDA graphs, rounded fused operators, split-KV attention,
phrase-tree verification, and packed recurrent KV storage. Keep optimization
switches disabled by default and preserve the eager training path.

Add complete R2R/RxR replay timing, a benchmark HTTP adapter, task-success
comparison semantics and documented navigation evidence.
Port the implementation to the current embodiinfer package namespace; retain
archived performance results as measurements of their frozen source snapshots.

Validation: 89 targeted CPU tests passed (19 skipped); 12 operator tests
passed on physical GPU 1. Full CPU suite: 515 passed, 80 skipped, with the
same five namespace compatibility failures reproduced on unmodified
origin/main (458 passed, 70 skipped). Ruff and formatting checks pass.

Signed-off-by: yufoo1 <yufoo1.cs@gmail.com>
Preserve independent histories and transactional outputs across batched graph, split-KV, and phrase-tree inference. Add native vLLM replay benchmarks with matching E2E and complete-forward boundaries and concise measured results.

Signed-off-by: yufoo1 <yufoo1.cs@gmail.com>
Signed-off-by: yufoo1 <yufoo1.cs@gmail.com>
- add official RxR prompt/action space (30/60/90 degrees, default 30) to prompt, processor, policy and factory; parse through policy.parse_actions in the decoder, batched runtime and phrase-tree speculation
- benchmarks: action_space config key for both engines; tensor-batch lane extended from 1/2/4 to include batch 8
- local benchmark lane only; upstream optimization implementation is unchanged

Signed-off-by: hootandy321 <133196559+hootandy321@users.noreply.github.com>
- _finish_generation no longer clones a full row history every turn: per-row owned buffers grow by a 4096-token alignment and are topped up in place on the same memory chain
- ownership is tracked with weakrefs: a different live owner receives a fresh buffer, so earlier memories keep their committed histories
- new counters owned_kv_allocations / owned_kv_reuses (~90% reuse in the measured runs)
- verified: 54 CPU tests pass; token-exact regression R2R B8 2997/2997 and RxR B4 3879/3879; RxR B4 peak allocation down 2.1 GiB; R2R B8 27.62 samples/s (+0.9%), RxR B4 13.11 samples/s (+0.5%)
- RxR B8 still OOM on the 32 GB card (pool + per-row owned KV coexist); deferred to the multi-base attention work

Signed-off-by: hootandy321 <133196559+hootandy321@users.noreply.github.com>
ruff format --check flags one expression in batching_activevln.py and
one call in modeling_activevln.py on this base. Formatting only.

Signed-off-by: longxmas <1185267696@qq.com>
- full RTX 5090 evidence for the optimized-branch matrix (EmbodiInfer graphs/tree/split-KV vs vLLM modern 0.30 full replay)
- 25 optimized reports including the OOM fix runs and both RxR B8 failures; 12 first-matrix reports; run/serve logs; token-diff tool
- fixed-replay inference throughput only, not navigation SR/SPL

Signed-off-by: hootandy321 <133196559+hootandy321@users.noreply.github.com>
The phrase tree (tree_decode) drafts a few common phrases, advances one
token per forward once a reply leaves the tree, and supports greedy
decoding only. With grammar_tree_budget=N the batched runtime drafts the
whole reply instead: a trie over the R2R/RxR action grammar, ranked by an
online prior built from the policy's own past replies, of which the N
most likely nodes are verified in one tree-attention forward.

Acceptance is exact. Greedy decoding keeps the argmax path and sampling
uses standard rejection sampling against the draft, so the output
distribution is unchanged and the grammar tree also works with
do_sample. The phrase tree still rejects sampling.

R2R, 64 episodes, B=4 greedy: decode forwards per batch drop from 2.95
with the phrase tree to 2.22 at N=32, and throughput rises from 19.75 to
22.28 frames/s (10.32 without a tree). With the official sampling
settings the grammar tree reaches 22.28 frames/s, 2.23x the sampling
baseline.

Signed-off-by: longxmas <1185267696@qq.com>
Opt-in runtime options for ActiveVLN batched replay; the defaults keep
the previous path.

- gpu_preprocess: uint8 frames are resized and normalised on the GPU by
  a bit-exact port of the slow Qwen2-VL image processor (Pillow bicubic
  fixed point, float64 rescale, float32 normalise). Tokens and positions
  are computed once per prompt text and frame size.
- attention_kernel="gqa": split-KV Triton kernel that packs the query
  heads sharing a KV head into one program; probabilities stay
  FP32-faithful (probability_parts=3).
- skip_padded_queries: GQA programs whose rows are all padding do no KV
  work; live rows keep the exact partition and reduction order.
- incremental_kv_commit: owned per-row KV buffers copy only the new turn
  instead of the whole history.
- vision_kernel="segmented"/"fused": attention inside each window or
  image with FP32-faithful probabilities instead of a dense masked SDPA;
  "fused" adds a fused QKV split + vision RoPE kernel and the rounded
  RMSNorm/SwiGLU kernels.
- benchmark_batch: uint8_observations hands the exact uint8 source frame
  to the processor; timings split host and device preparation.

On R2R (64 episodes, 4089 frames, greedy, grammar tree budget 32)
gpu_preprocess, incremental_kv_commit, uint8 observations and
skip_padded_queries reproduce their reference actions frame for frame;
gqa and the vision kernels change BF16 reduction order only.

Signed-off-by: longxmas <1185267696@qq.com>
Each decode step now makes one pinned non-blocking upload, memoizes stop
reasons, builds history pads from seen tokens, keeps scores as float
lists, compacts the pooled KV in batches and reuses the commit layer
views. Replays are bitwise identical (4089/4089 at B=4 and B=8).
benchmark_batch also accepts batch sizes 12 and 16.

Signed-off-by: longxmas <1185267696@qq.com>
Prefill padded every row of a batch to the longest turn: 18.8% of the
text-prefill tokens at B=8. With packed_prefill, the turns of a batch run
as one row of back-to-back tokens, padded only to a token bucket.

- gqa_split_kv_attention gains a packed mode: a program table of (row,
  query block) entries, per-row token starts and a device query extent.
  The extent equals the padded query length, so split partitions and
  results match the padded launch bit for bit.
- _text_forward writes packed KV with index_copy_ into per-token pool
  slots; bucket-padding tokens go to spare slots past the row capacities.
- Graphs are keyed by (token bucket, key bound per split count) and share
  static input buffers.

R2R replay (4089 frames, same-time baselines): B=8 N=23 39.56 -> 41.97
fps, B=4 N=31 34.48 -> 35.92 fps. Per-call divergence 7.1% / 7.5%, within
the batch-size noise floor (7.9-11.1%); no fallbacks; peak memory lower.

Signed-off-by: longxmas <1185267696@qq.com>
With gpu_preprocess, tokens and positions are cached per prompt, but each
cache miss (two per episode) still ran the Qwen processor's CPU image
pipeline only to discard its patches: 7.0 ms per miss.

- ActiveVLNProcessor.process_turn_tokens caches the image grid per frame
  size and expands and tokenizes the text exactly as the Qwen processor
  does: 0.85 ms per miss, identical fields on all 128 R2R turns checked.
- prepare_devices moves every row's token tensors in one pinned,
  non-blocking transfer instead of four synchronous copies per row;
  frames go through a pinned staging buffer.

R2R replay outputs are identical frame for frame (4089/4089). Preparation
drops from 0.85 to 0.64 ms/frame at B=8 and from 1.09 to 0.71 at B=4
(41.78 -> 42.36 fps at B=8 N=23).

Signed-off-by: longxmas <1185267696@qq.com>
With query_lengths, keys at or past position + count are masked for every
live query, yet the GQA kernel still loaded their values, so a NaN left in
those slots reached the output as 0 * NaN. Packed prefill writes its
padding tokens, whose outputs are unspecified, to spare slots that a pool
relocation can later hand to a row; the row then emits all-zero tokens
from the step its context reaches them, and so does the next episode in
that row.

Partitions still divide the full extent, so live outputs are bitwise
unchanged whenever the stale slots are finite. The failure surfaced on
R2R B=8 N=23 (episode 64 from step 83) and B=4 N=31 (episode 48 from
step 9) only after an unrelated change to the projection code.

Signed-off-by: longxmas <1185267696@qq.com>
Tree verification and singleton decode run every linear on 24-192 rows,
where cuBLAS leaves 10-35% of the 4090's bandwidth unused. With
decode_linear=True those projections use Triton kernels that accumulate
in FP32 in sequential K order without split-K, so results depend on
neither tiles nor row count and differ from cuBLAS only by output
rounding. Q, K and V share one launch; gate and up share one with the
rounded SwiGLU in the epilogue, bitwise equal to rounded_swiglu of the
two projections. Prefill row counts and narrow weights stay on cuBLAS.

R2R replay, on vs off on the same GPU (two GPUs each): B=8 N=23 +1.8% /
+2.0% (best 43.34 fps), B=4 N=31 +3.3% / +2.0% (best 37.58 fps), B=8
N=15 +3.0% / +2.0% (best 43.62 fps); decode drops 0.4-0.8 ms/frame.
Per-call divergence is 5.5-7.7%, below the batch-size noise floor
(7.9-11.1%). With the option off, outputs are identical frame for frame
to the parent commit.

Signed-off-by: longxmas <1185267696@qq.com>
perf(activevln): grammar-trie speculation and lossless throughput

Signed-off-by: Xingyu Liu <133196559+hootandy321@users.noreply.github.com>
Signed-off-by: yufoo1 <yufoo1.cs@gmail.com>
@yufoo1
yufoo1 force-pushed the bench/activevln-rxr-oom-fix branch from 5e968a9 to d7325b4 Compare October 21, 2026 10:40
Integrate speculative decoding, GPU preprocessing, CUDA graphs and
shared fused operators with the model-only benchmark services.
Publish portable 4090 and Orin NX profiles and verified first-128
navigation summaries from the recorded frozen sources.

Keep robot deployment and asynchronous execution-prefix changes in
their existing branches. The integrated tree has CPU validation;
recorded GPU measurements retain their original source provenance.

Signed-off-by: yufoo1 <yufoo1.cs@gmail.com>
Preserve the existing ActiveVLN review branch ancestry while keeping
the reviewed model-only tree based on current main. Its inference
changes are already integrated; raw lane-data logs and snapshot
archives remain excluded from the final tree.

Signed-off-by: yufoo1 <yufoo1.cs@gmail.com>
@yufoo1 yufoo1 changed the title perf(activevln): add batched inference optimizations perf(vln): optimize ActiveVLN and StreamVLN inference Oct 28, 2026
Keep device reports focused on the measured table, experimental conditions and
reproduction commands.

Signed-off-by: yufoo1 <yufoo1.cs@gmail.com>
@yufoo1
yufoo1 force-pushed the bench/activevln-rxr-oom-fix branch 2 times, most recently from bad7d4c to 603c87a Compare October 28, 2026 04:34
Signed-off-by: yufoo1 <yufoo1.cs@gmail.com>
@ZhouAo-ZA
ZhouAo-ZA merged commit d8cc710 into main Oct 28, 2026
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants