[Feature] Add opt-in KV-cache continuation for DualVLN evaluation

internnav/habitat_extensions/vln/habitat_vln_evaluator.py

@@ -261,6 +261,7 @@ def resume_from_output_path(self) -> None:
 
     def _run_eval_dual_system(self) -> tuple:  # noqa: C901
         self.model.eval()
+        kv_cache_continuation = getattr(self.model_args, 'kv_cache_continuation', False)
 
         # resume from previous results
         sucs, spls, oss, nes, ndtw = self.resume_from_output_path()
@@ -415,14 +416,17 @@ def _run_eval_dual_system(self) -> tuple:  # noqa: C901
                     inputs = self.processor(text=[text], images=input_images, return_tensors="pt").to(self.model.device)
 
                     with torch.no_grad():
-                        output_ids = self.model.generate(
+                        generation_outputs = self.model.generate(
                             **inputs,
                             max_new_tokens=128,
                             do_sample=False,
                             use_cache=True,
                             past_key_values=None,
                             return_dict_in_generate=True,
-                        ).sequences
+                        )
+                        output_ids = generation_outputs.sequences
+                        if not kv_cache_continuation:
+                            generation_outputs = None
 
                     llm_outputs = self.processor.tokenizer.decode(
                         output_ids[0][inputs.input_ids.shape[1] :], skip_special_tokens=True
@@ -445,7 +449,13 @@ def _run_eval_dual_system(self) -> tuple:  # noqa: C901
                         image_grid_thw = torch.cat([thw.unsqueeze(0) for thw in inputs.image_grid_thw], dim=0)
 
                         with torch.no_grad():
-                            traj_latents = self.model.generate_latents(output_ids, pixel_values, image_grid_thw)
+                            if kv_cache_continuation:
+                                traj_latents = self.model.generate_latents_from_cache(
+                                    generation_outputs, image_grid_thw, inputs.attention_mask
+                                )
+                            else:
+                                traj_latents = self.model.generate_latents(output_ids, pixel_values, image_grid_thw)
+                        generation_outputs = None
 
                         # prepocess align with navdp
                         image_dp = torch.tensor(np.array(look_down_image.resize((224, 224)))).to(torch.bfloat16) / 255
@@ -478,6 +488,7 @@ def _run_eval_dual_system(self) -> tuple:  # noqa: C901
                         print('predicted goal', pixel_goal, flush=True)
 
                     else:
+                        generation_outputs = None
                         action_seq = self.parse_actions(llm_outputs)
                         print('actions', action_seq, flush=True)
 

internnav/model/basemodel/internvla_n1/internvla_n1.py

@@ -7,6 +7,7 @@
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils.torch_utils import randn_tensor
 from transformers import (
+    DynamicCache,
     Qwen2_5_VLConfig,
     Qwen2_5_VLForConditionalGeneration,
     Qwen2_5_VLModel,
@@ -346,6 +347,69 @@ def generate_latents(self, input_ids, pixel_values, image_grid_thw):
 
         return hidden_states
 
+    @torch.no_grad()
+    def generate_latents_from_cache(self, generation_outputs, image_grid_thw, attention_mask=None):
+        """Read latent queries by continuing the cache from the same generation.
+
+        Use an unpadded, batch-one, non-beam generation's DynamicCache and its
+        original image grid/mask. The final token and latent queries are processed
+        together. The input cache is restored after readout. BF16 output may
+        differ from full replay because of floating-point execution order.
+        """
+        input_ids = generation_outputs.sequences
+        cache = getattr(generation_outputs, 'past_key_values', None)
+        if self.training:
+            raise ValueError('KV-cache latent readout requires model.eval().')
+        if type(cache) is not DynamicCache:
+            raise ValueError('KV-cache latent readout requires a DynamicCache from generate(use_cache=True).')
+        if (
+            input_ids.ndim != 2
+            or input_ids.shape[0] != 1
+            or getattr(generation_outputs, 'beam_indices', None) is not None
+        ):
+            raise ValueError('KV-cache latent readout supports batch size 1 and num_beams=1 only.')
+        if self.config.use_sliding_window:
+            raise ValueError('KV-cache latent readout requires full attention.')
+        if attention_mask is not None and (
+            attention_mask.ndim != 2
+            or attention_mask.shape[0] != 1
+            or not 0 < attention_mask.shape[1] <= input_ids.shape[1]
+            or not torch.all(attention_mask == 1)
+        ):
+            raise ValueError('KV-cache latent readout supports unpadded inputs only.')
+
+        cached_length = cache.get_seq_length()
+        sequence_length = input_ids.shape[1]
+        if cached_length <= 0 or sequence_length - cached_length not in (0, 1):
+            raise ValueError('Generation cache must cover the sequence except, optionally, its final token.')
+        if cache.key_cache[0].shape[0] != 1:
+            raise ValueError('KV-cache latent readout supports batch size 1 and num_beams=1 only.')
+        missing_ids = input_ids[:, cached_length:]
+        for token_id in (IMAGE_TOKEN_INDEX, self.config.video_token_id, TRAJ_TOKEN_INDEX):
+            if torch.any(missing_ids == token_id):
+                raise ValueError('The uncached generation suffix must contain ordinary text tokens only.')
+
+        n_query = self.get_n_query()
+        query_ids = input_ids.new_full((1, n_query), TRAJ_TOKEN_INDEX)
+        # Match full replay's mRoPE positions without depending on mutable rope_deltas.
+        position_ids, _ = self.get_rope_index(torch.cat([input_ids, query_ids], dim=1), image_grid_thw)
+        inputs_embeds = torch.cat([self.model.embed_tokens(missing_ids), self.model.latent_queries], dim=1)
+        cache_position = torch.arange(cached_length, sequence_length + n_query, device=input_ids.device)
+        try:
+            outputs = self.model(
+                inputs_embeds=inputs_embeds,
+                position_ids=position_ids[:, :, cached_length:],
+                past_key_values=cache,
+                cache_position=cache_position,
+                use_cache=False,
+                output_hidden_states=False,
+                return_dict=True,
+            )
+            return outputs.last_hidden_state[:, -n_query:, :]
+        finally:
+            # DynamicCache is updated in place, even with use_cache=False.
+            cache.crop(cached_length)
+
     def generate_traj(
         self,
         traj_latents,

scripts/eval/bash/eval_dual_system.sh

@@ -10,4 +10,5 @@ srun -p <YOUR_PARTITION_NAME> \
     --kill-on-bad-exit=1 \
     python scripts/eval/eval.py \
         --config $CONFIG \
+        "$@" \
     > logs/${MID_RUN_NAME}_log.txt 2>&1

scripts/eval/eval.py

@@ -19,6 +19,11 @@ def parse_args():
         default='scripts/eval/configs/h1_rdp_cfg.py',
         help='eval config file path, e.g. scripts/eval/configs/h1_cma_cfg.py',
     )
+    parser.add_argument(
+        '--kv-cache-continuation',
+        action='store_true',
+        help='Reuse the generation KV cache for latent readout in Habitat dual-system evaluation (default: off).',
+    )
     return parser.parse_args()
 
 
@@ -33,6 +38,10 @@ def load_eval_cfg(config_path, attr_name='eval_cfg'):
 def main():
     args = parse_args()
     evaluator_cfg = load_eval_cfg(args.config, attr_name='eval_cfg')
+    if args.kv_cache_continuation:
+        if evaluator_cfg.eval_type != 'habitat_vln' or evaluator_cfg.agent.model_settings.get('mode') != 'dual_system':
+            raise ValueError('--kv-cache-continuation requires a Habitat dual_system configuration.')
+        evaluator_cfg.agent.model_settings['kv_cache_continuation'] = True
 
     # fill in evaluator default config
     if evaluator_cfg.eval_type == 'vln_distributed':

tests/unit_test/kv_cache_continuation.md

@@ -0,0 +1,126 @@
+# KV-cache continuation: validation and scope
+
+The existing Habitat DualVLN entry point accepts one optional flag:
+
+```bash
+PYTHONPATH=. python scripts/eval/eval.py \
+  --config scripts/eval/configs/habitat_dual_system_cfg.py \
+  --kv-cache-continuation
+```
+
+Use the normal checkpoint/data setup. Omitting the flag keeps full replay.
+`scripts/eval/bash/eval_dual_system.sh` also forwards this flag; its existing
+Slurm setup is still required. No default configuration or dependency changes.
+
+## Automated regression checks
+
+```bash
+PYTHONPATH=. USE_TF=0 CUDA_VISIBLE_DEVICES='' OMP_NUM_THREADS=1 \
+  python -m pytest -q tests/unit_test/test_kv_cache_continuation.py
+```
+
+23 CPU tests passed without checkpoints or a simulator. The CLI tests exercise
+real parsing/configuration, replacing only evaluator initialization: default
+configuration unchanged, explicit opt-in, and rejection of unsupported mode.
+Tiny-model FP32 tests cover eager/SDPA, one/multiple images, EOS, complete or
+missing-final-token caches, stale position state, repeated readout, cache-prefix
+preservation, invalid inputs, skipped vision/logits and exception cleanup.
+Numerical tolerances are `atol=2e-6, rtol=2e-5`.
+
+The supported cache path is eval-only, unpadded batch one, non-beam generation,
+standard `DynamicCache`, full attention, and image input from the same generation.
+It processes the uncached final text token and four latent queries in one forward
+pass and restores the cache afterward. The original `generate_latents` is unchanged.
+
+## Habitat measurement
+
+Measured on 2026-09-05 against upstream main `7a5c62400ac45b313d9b709c740b64191556a242`
+plus this implementation. The submission base is dev
+`5e287ed395f0ad4996b7839e9ccd0f48c36870c4`; the DualVLN evaluation, model,
+configuration and dependency files used here are identical between those bases.
+CPU regressions were also rerun after changing the submission base.
+
+- GPU: one NVIDIA RTX A6000, 48 GB; BF16, FlashAttention 2.7.4.post1.
+- Python 3.9, PyTorch 2.6.0+cu124, Transformers 4.51.0, Habitat-Sim 0.2.4.
+- Local `InternVLA-N1` checkpoint; complete vision/language/System1 loading,
+  with no missing, unexpected or mismatched keys. No quantization or retraining.
+- Real R2R `val_unseen`, scene `2azQ1b91cZZ`, fixed episodes 10 and 11.
+  Python/NumPy/Torch/Habitat seed `100 + episode_id` at each reset.
+- Normal evaluator, observations, prompts and actions. History size 8;
+  front/history images 384 x 384, look-down image 640 x 480.
+  Original `nextdit_async` System1: 10 diffusion steps, 32 samples.
+- Only local paths, episode selection, seeding and the step cap were overridden
+  in memory. `max_steps_per_episode=40`; the existing loop reports 41 steps
+  for each episode. Auxiliary camera actions are counted separately below.
+
+### Paired latent readout (not the whole evaluation)
+
+At each of the seven actual pixel-goal calls, compare the two readouts using the
+same generation. Warm each path once, then time three repetitions per path in
+alternating order, with CUDA synchronization before and after each readout.
+Take each call's median, then summarize these seven medians. All seven calls
+have 10 images and one uncached final token. Generation, preprocessing, System1
+and simulator time are excluded; continuation includes final-token catch-up.
+
+| Episode / generation index | Prompt tokens | Full replay (ms) | Continuation (ms) |
+| --- | ---: | ---: | ---: |
+| 10 / 3 | 2297 | 553.608 | 31.095 |
+| 10 / 8 | 2295 | 556.985 | 31.399 |
+| 10 / 12 | 2296 | 551.997 | 31.116 |
+| 11 / 16 | 2272 | 529.133 | 31.177 |
+| 11 / 18 | 2275 | 531.158 | 31.209 |
+| 11 / 23 | 2276 | 529.865 | 31.308 |
+| 11 / 25 | 2276 | 529.031 | 31.275 |
+| Mean | | 540.25 | 31.23 |
+| Median | | 531.16 | 31.21 |
+
+Ratio of means: **17.30x latent-readout speedup**. Generation indices are
+zero-based across both episodes, including text-action calls.
+
+Mean per-query cosine similarity is 0.999700; minimum query cosine is 0.998746.
+Mean relative L2 error is 0.02576 (maximum 0.03628); maximum absolute latent
+difference is 2.0. These BF16 outputs are close, **not bitwise equal**.
+With the same System1 RNG state, all seven resulting discrete action lists match.
+Continuous trajectory values are not identical (maximum raw-output difference
+0.277344). Action comparison uses cloned trajectories because `traj_to_actions`
+modifies its input in place; shadow calls restore RNG state before navigation resumes.
+
+### Separate bounded closed-loop runs
+
+Run the same two episodes once per mode through the normal eval entry point,
+without paired/shadow model calls. Only continuation mode adds the flag.
+Measure the evaluator wall time with CUDA synchronization at its boundaries.
+This includes episode resets, environment steps, preprocessing, generation and
+System1, but excludes evaluator/model initialization. Per-stage synchronization
+adds instrumentation overhead. This is a small integration check, not a repeated
+statistical end-to-end benchmark.
+
+| Metric | Full replay | Continuation |
+| --- | ---: | ---: |
+| Two-episode evaluation time | 39.63 s | 35.91 s |
+| Generation / latent-readout calls | 26 / 7 | 26 / 7 |
+| Actual Habitat actions (including camera probes) | 433 | 433 |
+| SR / SPL / oracle success under this cap | 0 / 0 / 0 | 0 / 0 / 0 |
+| Mean navigation error | 6.65148 m | 6.65148 m |
+
+Observed wall-time ratio: **1.10x**, or 9.39% less time. All 433 actions match
+exactly (215 in episode 10, 218 in episode 11). The paired shadow run also
+preserves the baseline action stream. Latent readout accounted for approximately
+9.5% of baseline wall time, so its 17.30x local improvement does not imply 17.30x
+faster navigation.
+
+## Limits and reproducibility
+
+Repeat the unit command for inexpensive correctness checks. For Habitat A/B,
+use the same checkpoint, software/hardware, fixed episode subset and seeds,
+and compare full action streams, not just aggregate metrics. For paired timing,
+retain `generate(..., return_dict_in_generate=True, use_cache=True)` and compare
+`generate_latents(output.sequences, inputs.pixel_values, inputs.image_grid_thw)`
+with `generate_latents_from_cache(output, inputs.image_grid_thw, inputs.attention_mask)`
+before releasing that generation's cache, using the protocol above. Timing and
+episode-selection instrumentation was local; no benchmark framework is added.
+
+Both bounded episodes hit the step cap without success. Identical actions here
+do not establish unchanged full-dataset SR/SPL. Neither a complete navigation
+benchmark, robot deployment nor full upstream CI was validated. The default-off
+scope is intentional; BF16 numerical drift can affect other trajectories.