Fix DeMemWM noise-route pruning review findings

.exp_artifact/dememwm_noise_route_inference_pruning/diagnostic.py

@@ -25,7 +25,9 @@ STREAM_NOISE_BASE = {"anchor": 200, "dynamic": 300, "revisit": 400}
 BATCH_SIZE = 2
 TARGET_LENGTH = 3
 TIMING_WARMUP = 20
-TIMING_ITERS = 1000
+TIMING_ITERS = 200
+PROXY_PATCHES_PER_FRAME = 8
+PROXY_WIDTH = 32
 
 
 def _make_inputs(stream_masks):
@@ -89,8 +91,116 @@ def _target_proxy(packed_latents, frame_memory_segments, frame_memory_masks):
     return target + contribution.unsqueeze(0)
 
 
+def _make_proxy_weights(tensor):
+    dtype = tensor.dtype
+    device = tensor.device
+    in_dim = tensor.shape[-1] + 4 + 3 + 1
+    frame_proj = torch.linspace(
+        -0.17,
+        0.19,
+        steps=in_dim * PROXY_WIDTH,
+        dtype=dtype,
+        device=device,
+    ).reshape(in_dim, PROXY_WIDTH)
+    patch_embed = torch.linspace(-0.11, 0.13, steps=PROXY_PATCHES_PER_FRAME * PROXY_WIDTH, dtype=dtype, device=device).reshape(
+        PROXY_PATCHES_PER_FRAME,
+        PROXY_WIDTH,
+    )
+    q_proj = torch.linspace(-0.09, 0.10, steps=PROXY_WIDTH * PROXY_WIDTH, dtype=dtype, device=device).reshape(
+        PROXY_WIDTH,
+        PROXY_WIDTH,
+    )
+    k_proj = torch.linspace(0.08, -0.07, steps=PROXY_WIDTH * PROXY_WIDTH, dtype=dtype, device=device).reshape(
+        PROXY_WIDTH,
+        PROXY_WIDTH,
+    )
+    v_proj = torch.linspace(-0.05, 0.06, steps=PROXY_WIDTH * PROXY_WIDTH, dtype=dtype, device=device).reshape(
+        PROXY_WIDTH,
+        PROXY_WIDTH,
+    )
+    out_proj = torch.linspace(0.04, -0.03, steps=PROXY_WIDTH * PROXY_WIDTH, dtype=dtype, device=device).reshape(
+        PROXY_WIDTH,
+        PROXY_WIDTH,
+    )
+    return {
+        "frame_proj": frame_proj,
+        "patch_embed": patch_embed,
+        "q_proj": q_proj,
+        "k_proj": k_proj,
+        "v_proj": v_proj,
+        "out_proj": out_proj,
+    }
+
+
+def _packed_frame_valid_mask(frame_memory_segments, frame_memory_masks):
+    chunks = []
+    for key in ("target", *_DEMEMWM_STREAM_KEYS):
+        length = frame_memory_segments[key]
+        if length == 0:
+            continue
+        chunks.append(frame_memory_masks[key])
+    return torch.cat(chunks, dim=1)
+
+
+def _expand_frame_tokens(frame_tokens, frame_valid, patch_embed):
+    batch, frame_count, width = frame_tokens.shape
+    patch_tokens = frame_tokens[:, :, None, :] + patch_embed[None, None, :, :]
+    patch_tokens = patch_tokens.reshape(batch, frame_count * PROXY_PATCHES_PER_FRAME, width)
+    patch_valid = frame_valid[:, :, None].expand(batch, frame_count, PROXY_PATCHES_PER_FRAME).reshape(batch, -1)
+    return patch_tokens, patch_valid
+
+
+def _masked_attention(query, key, value, key_valid):
+    logits = torch.matmul(query, key.transpose(-1, -2)) / (query.shape[-1] ** 0.5)
+    logits = logits.masked_fill(~key_valid[:, None, :], torch.finfo(logits.dtype).min)
+    weights = torch.softmax(logits, dim=-1)
+    weights = torch.where(key_valid[:, None, :], weights, torch.zeros_like(weights))
+    return torch.matmul(weights, value)
+
+
+def _dit_like_inference_proxy(
+    packed_latents,
+    packed_conditions,
+    packed_poses,
+    packed_frame_indices,
+    frame_memory_segments,
+    frame_memory_masks,
+    proxy_weights,
+):
+    # Small deterministic stand-in for one DiT sample_step: frame features become
+    # patch tokens, then packed temporal attention and target-to-memory reads run
+    # over the same canonical segment/mask contract as the real call.
+    frame_index_feat = packed_frame_indices.to(dtype=packed_latents.dtype).unsqueeze(-1) / 100.0
+    frame_features = torch.cat([packed_latents, packed_conditions, packed_poses, frame_index_feat], dim=-1)
+    frame_tokens = torch.matmul(frame_features, proxy_weights["frame_proj"]).permute(1, 0, 2)
+    frame_valid = _packed_frame_valid_mask(frame_memory_segments, frame_memory_masks)
+    patch_tokens, patch_valid = _expand_frame_tokens(frame_tokens, frame_valid, proxy_weights["patch_embed"])
+
+    query = torch.matmul(patch_tokens, proxy_weights["q_proj"])
+    key = torch.matmul(patch_tokens, proxy_weights["k_proj"])
+    value = torch.matmul(patch_tokens, proxy_weights["v_proj"])
+    temporal = _masked_attention(query, key, value, patch_valid)
+    temporal = torch.matmul(torch.tanh(temporal), proxy_weights["out_proj"])
+
+    target_patch_count = frame_memory_segments["target"] * PROXY_PATCHES_PER_FRAME
+    target_tokens = temporal[:, :target_patch_count]
+    cursor = target_patch_count
+    memory_reads = []
+    for key_name in _DEMEMWM_STREAM_KEYS:
+        stream_patch_count = frame_memory_segments[key_name] * PROXY_PATCHES_PER_FRAME
+        stream_tokens = temporal[:, cursor : cursor + stream_patch_count]
+        stream_valid = patch_valid[:, cursor : cursor + stream_patch_count]
+        cursor += stream_patch_count
+        if stream_patch_count == 0:
+            continue
+        memory_reads.append(_masked_attention(target_tokens, stream_tokens, stream_tokens, stream_valid))
+    if memory_reads:
+        target_tokens = target_tokens + torch.stack(memory_reads, dim=0).sum(dim=0)
+    return target_tokens.mean(dim=(1, 2))
+
+
 def _time_proxy_call(fn, timing_tensor):
-    backend = "cuda_proxy" if timing_tensor.is_cuda else "cpu_proxy"
+    backend = "cuda_dit_like_proxy" if timing_tensor.is_cuda else "cpu_dit_like_proxy"
     with torch.no_grad():
         for _ in range(TIMING_WARMUP):
             output = fn()
@@ -110,8 +220,26 @@ def _full_proxy_inputs(inputs, routed_masks):
         [inputs["target_latents"], *[inputs["stream_latents_by_key"][key] for key in _DEMEMWM_STREAM_KEYS]],
         dim=0,
     )
+    packed_conditions = torch.cat(
+        [
+            inputs["target_conditions"],
+            *[
+                torch.zeros((STREAM_LENGTHS[key], BATCH_SIZE, inputs["target_conditions"].shape[-1]))
+                for key in _DEMEMWM_STREAM_KEYS
+            ],
+        ],
+        dim=0,
+    )
+    packed_poses = torch.cat(
+        [inputs["target_poses"], *[inputs["stream_poses_by_key"][key] for key in _DEMEMWM_STREAM_KEYS]],
+        dim=0,
+    )
+    packed_frame_indices = torch.cat(
+        [inputs["target_frame_indices"], *[inputs["stream_frame_indices_by_key"][key] for key in _DEMEMWM_STREAM_KEYS]],
+        dim=0,
+    )
     segments = {"target": TARGET_LENGTH, **{key: STREAM_LENGTHS[key] for key in _DEMEMWM_STREAM_KEYS}}
-    return packed_latents, segments, {"target": inputs["target_mask"], **routed_masks}
+    return packed_latents, packed_conditions, packed_poses, packed_frame_indices, segments, {"target": inputs["target_mask"], **routed_masks}
 
 
 def _expected_noise(memory_noise_levels_by_key, active_streams):
@@ -200,18 +328,45 @@ def _run_case(name, sampler_index, stream_masks, expected_active, expected_prune
         f"{name}: denoising noise-level shape contract mismatch",
     )
 
-    full_latents, full_proxy_segments, full_proxy_masks = _full_proxy_inputs(inputs, routed_masks)
-    full_output, full_step_ms, full_backend = _time_proxy_call(
-        lambda: _target_proxy(full_latents, full_proxy_segments, full_proxy_masks),
+    (
+        full_latents,
+        full_conditions,
+        full_poses,
+        full_frame_indices,
+        full_proxy_segments,
+        full_proxy_masks,
+    ) = _full_proxy_inputs(inputs, routed_masks)
+    full_target_output = _target_proxy(full_latents, full_proxy_segments, full_proxy_masks)
+    pruned_target_output = _target_proxy(active_packed_latents, active_segments, active_masks)
+    target_max_abs_diff = (full_target_output - pruned_target_output).abs().max().item()
+    _check(target_max_abs_diff == 0.0, f"{name}: proxy target output changed by {target_max_abs_diff}")
+
+    proxy_weights = _make_proxy_weights(full_latents)
+    _, full_step_ms, full_backend = _time_proxy_call(
+        lambda: _dit_like_inference_proxy(
+            full_latents,
+            full_conditions,
+            full_poses,
+            full_frame_indices,
+            full_proxy_segments,
+            full_proxy_masks,
+            proxy_weights,
+        ),
         full_latents,
     )
-    pruned_output, pruned_step_ms, pruned_backend = _time_proxy_call(
-        lambda: _target_proxy(active_packed_latents, active_segments, active_masks),
+    _, pruned_step_ms, pruned_backend = _time_proxy_call(
+        lambda: _dit_like_inference_proxy(
+            active_packed_latents,
+            active_packed_conditions,
+            active_frame_memory_pose,
+            active_frame_indices,
+            active_segments,
+            active_masks,
+            proxy_weights,
+        ),
         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")
 
@@ -219,6 +374,8 @@ def _run_case(name, sampler_index, stream_masks, expected_active, expected_prune
     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("  timing_scope=synthetic_dit_like_proxy_not_real_sample_step")
+    print("  speedup_ratio_scope=synthetic_proxy_only_not_measured_inference_speed")
     print(f"  full_tokens={full_tokens}")
     print(f"  pruned_tokens={pruned_tokens}")
     print(f"  token_reduction_pct={token_reduction_pct:.2f}")