Skip to content

perf(activevln): grammar-trie speculation and lossless throughput - #11

Merged
hootandy321 merged 8 commits into
bench/activevln-rxr-oom-fixfrom
perf/activevln-grammar-trie-throughput
Oct 21, 2026
Merged

perf(activevln): grammar-trie speculation and lossless throughput#11
hootandy321 merged 8 commits into
bench/activevln-rxr-oom-fixfrom
perf/activevln-grammar-trie-throughput

Conversation

@Longxmas

@Longxmas Longxmas commented Oct 20, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add grammar-trie draft trees for ActiveVLN with exact SpecInfer-style acceptance (greedy and sampling).
  • Stack lossless throughput work on top: GQA / vision kernels, GPU preprocess, packed prefill, decode host / upload cuts, and batch-invariant Triton decode GEMMs.
  • On R2R greedy, end-to-end fps rises from 10.3 (no tree) → 22.0 (grammar trie) → 37.0 at B=4, and to 43.2 at B=8; divergence stays in the same range as changing batch size alone.

Test plan

  • pytest ActiveVLN / GQA / vision / decode_linear suites on the lab 4090 env
  • Replay R2R B=4 and B=8 with the final config (attention_kernel=gqa, gpu_preprocess, packed_prefill, decode_linear, …)
  • Spot-check greedy and sampling outputs against the grammar-trie baseline for per-call divergence vs batch-size noise

@Longxmas
style(activevln): apply ruff format to the batch runtime and model
32e354a 
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>
@Longxmas
feat(activevln): grammar-trie tree speculation with exact acceptance
4b728e2 
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>
@Longxmas
perf(activevln): lossless batched replay throughput options
938fe7b 
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>
@Longxmas
perf(activevln): cut grammar-tree decode host overhead
41dcd00 
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>
@Longxmas
perf(activevln): pack batched prefill turns back to back in CUDA graphs
a7b80f9 
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>
@Longxmas
perf(activevln): tokenize turns without the CPU image pipeline
2b3a84c 
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>
@Longxmas
fix(attention): never read KV slots past a row's live queries
e9133f4 
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>
@Longxmas
perf(activevln): batch-invariant Triton GEMMs for decode projections
1cd8c63 
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>
@yufoo1
yufoo1 self-requested a review October 20, 2026 17:46
@hootandy321
hootandy321 merged commit d288bad into bench/activevln-rxr-oom-fix Oct 21, 2026
1 of 3 checks passed
@hootandy321
hootandy321 deleted the perf/activevln-grammar-trie-throughput branch October 21, 2026 02:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@yufoo1 yufoo1 Awaiting requested review from yufoo1

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@Longxmas @hootandy321