From ee40c8acd57a1d3342633956b992c15148f148aa Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Wed, 8 Apr 2026 06:40:58 +0000
Subject: [PATCH] perf: optimize VLA token loop with torch.unbind

Extracts dtype casting and intermediate tensor slicing (`[:, t, :]`) out of the token loop. By using `torch.unbind(dim=1)` prior to the loop, it removes the per-step PyTorch view operations and dispatch overhead, replacing them with simple list indexing.

Co-authored-by: Vishal-sys-code <68536727+Vishal-sys-code@users.noreply.github.com>
---
 benchmark_vla.py            | 36 ++++++++++++++++++++++++++++++++++++
 benchmark_vla_advanced.py   | 36 ++++++++++++++++++++++++++++++++++++
 src/models/attention/vla.py | 17 ++++++++++++-----
 3 files changed, 84 insertions(+), 5 deletions(-)
 create mode 100644 benchmark_vla.py
 create mode 100644 benchmark_vla_advanced.py

diff --git a/benchmark_vla.py b/benchmark_vla.py
new file mode 100644
index 0000000..860b8c6
--- /dev/null
+++ b/benchmark_vla.py
@@ -0,0 +1,36 @@
+import time
+import torch
+from src.models.attention.vla import VLALayer
+
+def run_benchmark():
+    B = 16
+    T = 1000
+    d_model = 64
+    device = torch.device("cpu")
+
+    # Initialize the model
+    layer = VLALayer(d_model=d_model).to(device)
+
+    # Create input tensor
+    x = torch.randn(B, T, d_model, device=device)
+
+    # Warmup
+    print("Warming up...")
+    with torch.no_grad():
+        for _ in range(3):
+            _ = layer(x)
+
+    # Benchmark
+    print("Benchmarking...")
+    num_iters = 10
+    start_time = time.time()
+    with torch.no_grad():
+        for _ in range(num_iters):
+            _ = layer(x)
+    end_time = time.time()
+
+    avg_time = (end_time - start_time) / num_iters
+    print(f"Average time per forward pass: {avg_time:.4f} seconds")
+
+if __name__ == "__main__":
+    run_benchmark()
diff --git a/benchmark_vla_advanced.py b/benchmark_vla_advanced.py
new file mode 100644
index 0000000..7a5d051
--- /dev/null
+++ b/benchmark_vla_advanced.py
@@ -0,0 +1,36 @@
+import time
+import torch
+from src.models.attention.vla import VLALayer
+
+def run_benchmark():
+    B = 16
+    T = 2000 # Increased length for clearer effect
+    d_model = 64
+    device = torch.device("cpu")
+
+    # Initialize the model
+    layer = VLALayer(d_model=d_model).to(device)
+
+    # Create input tensor
+    x = torch.randn(B, T, d_model, device=device)
+
+    # Warmup
+    print("Warming up...")
+    with torch.no_grad():
+        for _ in range(5):
+            _ = layer(x)
+
+    # Benchmark
+    print("Benchmarking...")
+    num_iters = 20
+    start_time = time.time()
+    with torch.no_grad():
+        for _ in range(num_iters):
+            _ = layer(x)
+    end_time = time.time()
+
+    avg_time = (end_time - start_time) / num_iters
+    print(f"Average time per forward pass: {avg_time:.4f} seconds")
+
+if __name__ == "__main__":
+    run_benchmark()
diff --git a/src/models/attention/vla.py b/src/models/attention/vla.py
index e6bfe2f..3e0cb2e 100644
--- a/src/models/attention/vla.py
+++ b/src/models/attention/vla.py
@@ -123,19 +123,26 @@ def forward(self, x: torch.Tensor, return_states: bool = False, symbolic_adj: Op
         mem_step = self.memory_manager.step.item()
         fallback_count = self.inverse_tracker.fallback_count.item()
 
+        # Pre-cast and unbind sequence tensors to remove slicing overhead in loop
+        Q_list = Q.to(dtype=torch.float32).unbind(dim=1)
+        K_list = K.to(dtype=torch.float32).unbind(dim=1)
+        V_list = V.to(dtype=torch.float32).unbind(dim=1)
+        Lambda_list = Lambda_seq.to(dtype=torch.float32).unbind(dim=1)
+        U_list = U_seq.to(dtype=torch.float32).unbind(dim=1)
+
         # Iterate over tokens
         for t in range(T):
             # Extract current timestep pre-computed vectors
-            q_t = Q[:, t, :].to(dtype=torch.float32)  # (B, d_head)
-            k_t = K[:, t, :].to(dtype=torch.float32)  # (B, d_head)
-            v_t = V[:, t, :].to(dtype=torch.float32)  # (B, d_head)
+            q_t = Q_list[t]  # (B, d_head)
+            k_t = K_list[t]  # (B, d_head)
+            v_t = V_list[t]  # (B, d_head)
 
             # Step 4.2: Compute score s_t (legacy, unused for routing but kept for diagnostics if needed)
             s_t = (k_t * q_t).sum(dim=-1, keepdim=True)  # (B, 1)
 
             # Step 4.3: Extract pre-computed penalty components
-            lambda_t = Lambda_seq[:, t, :].to(dtype=torch.float32)
-            u_t = U_seq[:, t, :].to(dtype=torch.float32)
+            lambda_t = Lambda_list[t]
+            u_t = U_list[t]
             
             # Step 4.3b: Scale u_t to prevent large rank-1 updates
             import math