K2.A.1 fix: cast KakeyaLattice round-tripped K/V back to resident dtype

_round_trip_resident_through_compressor crashed on CUDA bf16:
  index_copy_(): self and source expected to have the same dtype,
  but got (self) BFloat16 and (source) Float

KakeyaLattice's compress/decompress runs its quantize/dequantize math in
fp32 for fidelity, so the round-tripped K/V come back fp32 while the
resident K/V cache is the model compute dtype (bf16 on CUDA). index_copy_
requires matching dtype+device. Cast K/V_round_tripped back to the
destination's dtype+device before writing. (Didn't surface on Mac where
the fp32 path matched.)

Co-authored-by: FluffyAIcode <FluffyAIcode@users.noreply.github.com>

Author: cursoragent

inference_engine/v04/dlm_restored_verifier.py

@@ -201,8 +201,17 @@ def _round_trip_resident_through_compressor(
     # resident positions get the round-tripped values.
     K_out = K.clone()
     V_out = V.clone()
-    K_out.index_copy_(-2, pos_tensor, K_round_tripped)
-    V_out.index_copy_(-2, pos_tensor, V_round_tripped)
+    # The compressor round-trip may upcast to fp32 (KakeyaLattice's
+    # quantize/dequantize math runs in fp32 for numerical fidelity),
+    # whereas the resident K/V cache is the model's compute dtype
+    # (bf16 on CUDA). index_copy_ requires matching dtype (and device),
+    # so cast the round-tripped tensors back before writing them in.
+    K_out.index_copy_(
+        -2, pos_tensor, K_round_tripped.to(device=K_out.device, dtype=K_out.dtype),
+    )
+    V_out.index_copy_(
+        -2, pos_tensor, V_round_tripped.to(device=V_out.device, dtype=V_out.dtype),
+    )
     return K_out, V_out