Measure DeMemWM inference noise-route pruning

.exp_artifact/dememwm_noise_route_inference_pruning/diagnostic.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 from types import SimpleNamespace
 import sys
+import time
 
 import torch
 
@@ -23,6 +24,8 @@ STREAM_VALUES = {"anchor": 20.0, "dynamic": 30.0, "revisit": 40.0}
 STREAM_NOISE_BASE = {"anchor": 200, "dynamic": 300, "revisit": 400}
 BATCH_SIZE = 2
 TARGET_LENGTH = 3
+TIMING_WARMUP = 20
+TIMING_ITERS = 1000
 
 
 def _make_inputs(stream_masks):
@@ -86,6 +89,22 @@ def _target_proxy(packed_latents, frame_memory_segments, frame_memory_masks):
     return target + contribution.unsqueeze(0)
 
 
+def _time_proxy_call(fn, timing_tensor):
+    backend = "cuda_proxy" if timing_tensor.is_cuda else "cpu_proxy"
+    with torch.no_grad():
+        for _ in range(TIMING_WARMUP):
+            output = fn()
+        if timing_tensor.is_cuda:
+            torch.cuda.synchronize(timing_tensor.device)
+        start = time.perf_counter()
+        for _ in range(TIMING_ITERS):
+            output = fn()
+        if timing_tensor.is_cuda:
+            torch.cuda.synchronize(timing_tensor.device)
+        elapsed_ms = (time.perf_counter() - start) * 1000.0 / TIMING_ITERS
+    return output, elapsed_ms, backend
+
+
 def _full_proxy_inputs(inputs, routed_masks):
     packed_latents = torch.cat(
         [inputs["target_latents"], *[inputs["stream_latents_by_key"][key] for key in _DEMEMWM_STREAM_KEYS]],
@@ -142,6 +161,8 @@ def _run_case(name, sampler_index, stream_masks, expected_active, expected_prune
     )
 
     full_segments = {"target": TARGET_LENGTH, **{key: STREAM_LENGTHS[key] for key in _DEMEMWM_STREAM_KEYS}}
+    full_tokens = sum(full_segments.values())
+    pruned_tokens = sum(active_segments.values())
     full_memory_tokens = sum(full_segments[key] for key in _DEMEMWM_STREAM_KEYS)
     pruned_memory_tokens = sum(active_segments[key] for key in _DEMEMWM_STREAM_KEYS)
 
@@ -156,7 +177,7 @@ def _run_case(name, sampler_index, stream_masks, expected_active, expected_prune
         _check(active_segments[key] == STREAM_LENGTHS[key], f"{name}: active {key} segment length changed")
         _check(torch.equal(active_masks[key], routed_masks[key]), f"{name}: active {key} mask changed")
 
-    expected_total_len = TARGET_LENGTH + pruned_memory_tokens
+    expected_total_len = pruned_tokens
     _check(active_packed_latents.shape == (expected_total_len, BATCH_SIZE, 2), f"{name}: packed latent shape mismatch")
     _check(active_packed_conditions.shape == (expected_total_len, BATCH_SIZE, 4), f"{name}: packed condition shape mismatch")
     _check(active_frame_memory_pose.shape == (expected_total_len, BATCH_SIZE, 3), f"{name}: packed pose shape mismatch")
@@ -180,14 +201,30 @@ def _run_case(name, sampler_index, stream_masks, expected_active, expected_prune
     )
 
     full_latents, full_proxy_segments, full_proxy_masks = _full_proxy_inputs(inputs, routed_masks)
-    full_output = _target_proxy(full_latents, full_proxy_segments, full_proxy_masks)
-    pruned_output = _target_proxy(active_packed_latents, active_segments, active_masks)
+    full_output, full_step_ms, full_backend = _time_proxy_call(
+        lambda: _target_proxy(full_latents, full_proxy_segments, full_proxy_masks),
+        full_latents,
+    )
+    pruned_output, pruned_step_ms, pruned_backend = _time_proxy_call(
+        lambda: _target_proxy(active_packed_latents, active_segments, active_masks),
+        active_packed_latents,
+    )
+    _check(full_backend == pruned_backend, f"{name}: timing backend changed across proxy calls")
     target_max_abs_diff = (full_output - pruned_output).abs().max().item()
     _check(target_max_abs_diff == 0.0, f"{name}: proxy target output changed by {target_max_abs_diff}")
+    token_reduction_pct = 100.0 * (1.0 - float(pruned_tokens) / float(full_tokens))
+    speedup_ratio = full_step_ms / pruned_step_ms if pruned_step_ms > 0.0 else float("inf")
 
     print(f"case={name}")
     print(f"  route=anchor:high,dynamic:low,revisit:all")
     print(f"  denoising_step={sampler_index} noise_level={_sampler_index_noise(sampler_index)}")
+    print(f"  timing_backend={full_backend} timing_iters={TIMING_ITERS}")
+    print(f"  full_tokens={full_tokens}")
+    print(f"  pruned_tokens={pruned_tokens}")
+    print(f"  token_reduction_pct={token_reduction_pct:.2f}")
+    print(f"  full_step_ms={full_step_ms:.6f}")
+    print(f"  pruned_step_ms={pruned_step_ms:.6f}")
+    print(f"  speedup_ratio={speedup_ratio:.3f}")
     print(f"  active_streams={active_streams}")
     print(f"  pruned_streams={pruned_streams}")
     print(f"  full_memory_token_count={full_memory_tokens}")