diff --git a/internnav/habitat_extensions/vln/habitat_vln_evaluator.py b/internnav/habitat_extensions/vln/habitat_vln_evaluator.py index 01378df3..e93f98e7 100644 --- a/internnav/habitat_extensions/vln/habitat_vln_evaluator.py +++ b/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) diff --git a/internnav/model/basemodel/internvla_n1/internvla_n1.py b/internnav/model/basemodel/internvla_n1/internvla_n1.py index 69c7f474..928133fe 100644 --- a/internnav/model/basemodel/internvla_n1/internvla_n1.py +++ b/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, diff --git a/scripts/eval/bash/eval_dual_system.sh b/scripts/eval/bash/eval_dual_system.sh index 91f17a54..84d014de 100755 --- a/scripts/eval/bash/eval_dual_system.sh +++ b/scripts/eval/bash/eval_dual_system.sh @@ -10,4 +10,5 @@ srun -p \ --kill-on-bad-exit=1 \ python scripts/eval/eval.py \ --config $CONFIG \ + "$@" \ > logs/${MID_RUN_NAME}_log.txt 2>&1 diff --git a/scripts/eval/eval.py b/scripts/eval/eval.py index eb6c17df..8b381738 100644 --- a/scripts/eval/eval.py +++ b/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': diff --git a/tests/unit_test/kv_cache_continuation.md b/tests/unit_test/kv_cache_continuation.md new file mode 100644 index 00000000..06849db6 --- /dev/null +++ b/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. diff --git a/tests/unit_test/test_kv_cache_continuation.py b/tests/unit_test/test_kv_cache_continuation.py new file mode 100644 index 00000000..02c3646e --- /dev/null +++ b/tests/unit_test/test_kv_cache_continuation.py @@ -0,0 +1,243 @@ +"""Opt-in CLI and tiny-model cache regressions; no checkpoints or simulator. + +See kv_cache_continuation.md for Habitat measurements and their limitations. +""" + +import copy +import importlib.util +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from internnav.model.basemodel.internvla_n1.internvla_n1 import ( + IMAGE_TOKEN_INDEX, + InternVLAN1ForCausalLM, + InternVLAN1ModelConfig, +) + + +@pytest.fixture +def entrypoint(monkeypatch): + # Only replace the simulator entry point; exercise the real argument parser + # and the repository's real configuration loading and override logic. + evaluator = ModuleType('internnav.evaluator') + evaluator.Evaluator = SimpleNamespace(init=Mock(return_value=SimpleNamespace(eval=Mock()))) + monkeypatch.setitem(sys.modules, 'internnav.evaluator', evaluator) + root = Path(__file__).resolve().parents[2] + spec = importlib.util.spec_from_file_location('eval_cli_test', root / 'scripts/eval/eval.py') + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module, root + + +@pytest.mark.parametrize('enabled', [False, True]) +def test_flag_only_changes_requested_setting(entrypoint, monkeypatch, enabled): + module, root = entrypoint + config = str(root / 'scripts/eval/configs/habitat_dual_system_cfg.py') + original = module.load_eval_cfg(config) + args = ['eval.py', '--config', config] + if enabled: + args.append('--kv-cache-continuation') + monkeypatch.setattr(sys, 'argv', args) + assert module.parse_args().kv_cache_continuation is enabled + module.main() + actual = module.Evaluator.init.call_args[0][0] + if enabled: + original.agent.model_settings['kv_cache_continuation'] = True + assert actual == original + module.Evaluator.init.return_value.eval.assert_called_once_with() + + +def test_rejects_flag_for_system2_only(entrypoint, monkeypatch): + module, root = entrypoint + config = str(root / 'scripts/eval/configs/habitat_s2_cfg.py') + monkeypatch.setattr(sys, 'argv', ['eval.py', '--config', config, '--kv-cache-continuation']) + with pytest.raises(ValueError, match='Habitat dual_system'): + module.main() + module.Evaluator.init.assert_not_called() + + +@pytest.fixture(scope='module', params=['eager', 'sdpa']) +def model(request): + config = InternVLAN1ModelConfig( + vocab_size=151680, + hidden_size=24, + intermediate_size=48, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + rope_scaling={'type': 'mrope', 'mrope_section': [2, 2, 2]}, + image_token_id=IMAGE_TOKEN_INDEX, + video_token_id=151656, + vision_start_token_id=151652, + vision_end_token_id=151653, + eos_token_id=151645, + pad_token_id=151643, + n_query=4, + system1='navdp', # Synchronous NavDP constructs only the latent queries. + vision_config={ + 'depth': 1, + 'hidden_size': 24, + 'intermediate_size': 48, + 'num_heads': 2, + 'out_hidden_size': 24, + 'patch_size': 14, + 'spatial_merge_size': 2, + 'temporal_patch_size': 2, + 'fullatt_block_indexes': [0], + }, + ) + config._attn_implementation = request.param + previous_threads = torch.get_num_threads() + try: + torch.set_num_threads(1) + with torch.random.fork_rng(devices=[]): + torch.random.default_generator.manual_seed(0) + yield InternVLAN1ForCausalLM(config).eval() + finally: + torch.set_num_threads(previous_threads) + + +def make_inputs(image_count): + grids = torch.tensor([[1, 4, 4], [1, 4, 6]][:image_count]) + ids = [30, 31] + for grid in grids: + ids += [151652] + [IMAGE_TOKEN_INDEX] * (int(grid.prod()) // 4) + [151653, 32] + input_ids = torch.tensor([ids + [33, 34]]) + return { + 'input_ids': input_ids, + 'attention_mask': torch.ones_like(input_ids), + 'pixel_values': torch.randn(int(grids.prod(dim=-1).sum()), 3 * 2 * 14 * 14), + 'image_grid_thw': grids, + } + + +@torch.no_grad() +def generate(model, inputs, eos=False): + return model.generate( + **inputs, + do_sample=False, + max_new_tokens=3, + eos_token_id=151645 if eos else None, + forced_eos_token_id=151645 if eos else None, + use_cache=True, + return_dict_in_generate=True, + ) + + +@pytest.mark.parametrize('image_count', [1, 2]) +@pytest.mark.parametrize('complete_cache', [False, True]) +@pytest.mark.parametrize('eos', [False, True]) +@torch.no_grad() +def test_matches_full_replay(model, image_count, complete_cache, eos): + inputs = make_inputs(image_count) + generated = generate(model, inputs, eos) + ids, cache = generated.sequences, generated.past_key_values + assert cache.get_seq_length() == ids.shape[1] - 1 + if eos: + assert ids[0, -1] == 151645 + if complete_cache: + positions, _ = model.get_rope_index(ids, inputs['image_grid_thw']) + model( + input_ids=ids[:, -1:], + position_ids=positions[:, :, -1:], + past_key_values=cache, + cache_position=torch.tensor([ids.shape[1] - 1]), + use_cache=True, + ) + original_length = cache.get_seq_length() + original_kv = [(key.clone(), value.clone()) for key, value in cache] + full = model.generate_latents(ids, inputs['pixel_values'], inputs['image_grid_thw']) + # An unrelated request may change this state before readout. + model.rope_deltas = torch.tensor([[12345]]) + for _ in range(2): + actual = model.generate_latents_from_cache(generated, inputs['image_grid_thw'], inputs['attention_mask']) + torch.testing.assert_close(actual, full, atol=2e-6, rtol=2e-5) + assert not actual.requires_grad + assert cache.get_seq_length() == original_length + for (key, value), (expected_key, expected_value) in zip(cache, original_kv): + torch.testing.assert_close(key, expected_key, atol=0, rtol=0) + torch.testing.assert_close(value, expected_value, atol=0, rtol=0) + + +@torch.no_grad() +def test_rejects_unsupported_inputs(model): + inputs = make_inputs(1) + generated = generate(model, inputs) + grid = inputs['image_grid_thw'] + with pytest.raises(ValueError, match='DynamicCache'): + model.generate_latents_from_cache(SimpleNamespace(sequences=generated.sequences), grid) + with pytest.raises(ValueError, match='unpadded'): + model.generate_latents_from_cache(generated, grid, torch.tensor([[0, 1]])) + with pytest.raises(ValueError, match='unpadded'): + model.generate_latents_from_cache(generated, grid, torch.ones(1, 1, 1, 1)) + with pytest.raises(ValueError, match='num_beams'): + model.generate_latents_from_cache(SimpleNamespace(**dict(generated), beam_indices=torch.tensor([0])), grid) + with pytest.raises(ValueError, match='batch size 1'): + model.generate_latents_from_cache( + SimpleNamespace(sequences=generated.sequences.repeat(2, 1), past_key_values=generated.past_key_values), + grid, + ) + bad = copy.deepcopy(generated) + bad.past_key_values.crop(0) + with pytest.raises(ValueError, match='final token'): + model.generate_latents_from_cache(bad, grid) + bad = copy.deepcopy(generated) + bad.past_key_values.crop(bad.past_key_values.get_seq_length() - 1) + with pytest.raises(ValueError, match='final token'): + model.generate_latents_from_cache(bad, grid) + with pytest.raises(ValueError, match='final token'): + model.generate_latents_from_cache( + SimpleNamespace(sequences=generated.sequences[:, :-2], past_key_values=generated.past_key_values), + grid, + ) + bad = copy.deepcopy(generated) + bad.sequences[0, -1] = IMAGE_TOKEN_INDEX + with pytest.raises(ValueError, match='ordinary text'): + model.generate_latents_from_cache(bad, grid) + model.config.use_sliding_window = True + try: + with pytest.raises(ValueError, match='full attention'): + model.generate_latents_from_cache(generated, grid) + finally: + model.config.use_sliding_window = False + model.train() + try: + with pytest.raises(ValueError, match='model.eval'): + model.generate_latents_from_cache(generated, grid) + finally: + model.eval() + + +@torch.no_grad() +def test_avoids_vision_and_logits_and_restores_cache_on_error(model): + inputs = make_inputs(1) + generated = generate(model, inputs) + cache = generated.past_key_values + original_length = cache.get_seq_length() + + def unexpected_forward(*args): + raise AssertionError('Readout must not run vision or the vocabulary head') + + hooks = [module.register_forward_pre_hook(unexpected_forward) for module in (model.visual, model.lm_head)] + try: + model.generate_latents_from_cache(generated, inputs['image_grid_thw']) + finally: + for hook in hooks: + hook.remove() + + def fail_after_forward(*args): + assert cache.get_seq_length() > original_length + raise RuntimeError('simulated failure after cache update') + + hook = model.model.register_forward_hook(fail_after_forward) + try: + with pytest.raises(RuntimeError, match='simulated failure'): + model.generate_latents_from_cache(generated, inputs['image_grid_thw']) + finally: + hook.remove() + assert cache.get_seq_length() == original_length