From 0b0ef8f638d28311b9daa468b75053879612f6fe Mon Sep 17 00:00:00 2001 From: "Vishal S. Pandey" Date: Sun, 12 Apr 2026 20:35:23 +0530 Subject: [PATCH 1/2] Update benchmark_retrieval.py --- scripts/benchmark_retrieval.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/scripts/benchmark_retrieval.py b/scripts/benchmark_retrieval.py index d4532b6..556fcde 100644 --- a/scripts/benchmark_retrieval.py +++ b/scripts/benchmark_retrieval.py @@ -9,6 +9,7 @@ from src.models.attention.vla import VLALayer from src.models.attention.deltanet import DeltaNetLayer from src.models.attention.linear_transformer import LinearTransformerLayer +from src.models.attention.fast_vla import VLAParallel, VLATriton class CausalConv1d(nn.Module): """ @@ -150,7 +151,7 @@ def train_and_eval(model, args, device): def main(): parser = argparse.ArgumentParser(description="Associative Retrieval Benchmark") - parser.add_argument("--model", type=str, required=True, choices=["VLA", "DeltaNet", "Linear"], help="Model architecture") + parser.add_argument("--model", type=str, required=True, choices=["VLA", "DeltaNet", "Linear", "VLAMamba", "VLATriton", "VLAKVExploding", "VLACombined"], help="Model architecture") parser.add_argument("--seq_len", type=int, required=True, help="Sequence Length") parser.add_argument("--epochs", type=int, default=100, help="Number of training epochs") parser.add_argument("--batch_size", type=int, default=16, help="Batch size") @@ -169,16 +170,25 @@ def main(): models = { "VLA": VLALayer, "DeltaNet": DeltaNetLayer, - "Linear": LinearTransformerLayer + "Linear": LinearTransformerLayer, + "VLAMamba": VLAParallel, + "VLATriton": VLATriton, + "VLAKVExploding": VLAParallel, + "VLACombined": VLATriton } layer_cls = models[args.model] - if args.model == "VLA": + if args.model in ["VLA", "VLAMamba", "VLATriton", "VLAKVExploding", "VLACombined"]: # Pass lambda_0=0.1 as requested by the user for stabilization + # and configure use_kv_exploding_fix + use_kv_exploding_fix = args.model in ["VLAKVExploding", "VLACombined"] + class VLALayerWithLambda(layer_cls): def __init__(self, *args_init, **kwargs_init): kwargs_init["lambda_0"] = 0.1 + if args.model != "VLA": + kwargs_init["use_kv_exploding_fix"] = use_kv_exploding_fix super().__init__(*args_init, **kwargs_init) model = ModelWrapper(VLALayerWithLambda, args.d_model, args.vocab_size).to(device) else: From 79f5ce0f502c89fa03ee3d20fa6d2e03779d7a9e Mon Sep 17 00:00:00 2001 From: "Vishal S. Pandey" Date: Sun, 12 Apr 2026 20:37:50 +0530 Subject: [PATCH 2/2] Implement fast VLA attention with Triton MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Author: @Akshay594 Add a new fast VLA attention module implementing sequential, parallel (prefix-scan) and an optional Triton-fused variant (src/models/attention/fast_vla.py). Include parallel-scan and Sherman–Morrison scan implementations, KV-explosion fixes, and Triton kernel guarded by HAS_TRITON. Add a basic unit test for the parallel variant (tests/test_fast_vla.py) and a speed benchmark script (scripts/benchmark_speed.py). Update plotting script (scripts/plot_retrieval.py) to aggregate retrieval CSVs and produce a combined accuracy plot. Scripts save figures under results/ and handle CPU/GPU/Triton availability. --- scripts/benchmark_speed.py | 88 ++++++++ scripts/plot_retrieval.py | 94 +++++---- src/models/attention/fast_vla.py | 344 +++++++++++++++++++++++++++++++ tests/test_fast_vla.py | 22 ++ 4 files changed, 505 insertions(+), 43 deletions(-) create mode 100644 scripts/benchmark_speed.py create mode 100644 src/models/attention/fast_vla.py create mode 100644 tests/test_fast_vla.py diff --git a/scripts/benchmark_speed.py b/scripts/benchmark_speed.py new file mode 100644 index 0000000..368ad47 --- /dev/null +++ b/scripts/benchmark_speed.py @@ -0,0 +1,88 @@ +import time +import torch +import matplotlib.pyplot as plt +from src.models.attention.fast_vla import VLASequential, VLAParallel, VLATriton, HAS_TRITON + +# Wong color palette (colorblind friendly) +WONG_PALETTE = { + "black": "#000000", + "orange": "#E69F00", + "skyblue": "#56B4E9", + "bluishgreen": "#009E73", + "yellow": "#F0E442", + "blue": "#0072B2", + "vermillion": "#D55E00", + "reddishpurple": "#CC79A7" +} + +def bench(layer, x, n=20): + layer.eval() + with torch.no_grad(): + for _ in range(5): + layer(x) + if x.device.type == "cuda": + torch.cuda.synchronize() + t0 = time.perf_counter() + with torch.no_grad(): + for _ in range(n): + layer(x) + if x.device.type == "cuda": + torch.cuda.synchronize() + return (time.perf_counter() - t0) / n * 1000 # ms + +def main(): + device = "cuda" if torch.cuda.is_available() else "cpu" + d, B = 64, 4 + + print(f"d={d} B={B} device={device}\n") + print(f"{'T':>6} {'Sequential':>12} {'Parallel':>10} {'Triton':>8} {'Tri/Seq':>8}") + print("-" * 56) + + seq_lens = [64, 256, 512, 1024, 2048, 4096] + if device == "cpu": + seq_lens = [64, 128, 256, 512] # Limit for CPU + + results = { + "Sequential": [], + "Parallel": [], + "Triton": [] + } + + for T in seq_lens: + x = torch.randn(B, T, d).to(device) + seq = VLASequential(d_model=d).to(device) + par = VLAParallel(d_model=d).to(device) + + t_seq = bench(seq, x) + t_par = bench(par, x) + + results["Sequential"].append(t_seq) + results["Parallel"].append(t_par) + + if HAS_TRITON and device == "cuda": + tri = VLATriton(d_model=d).to(device) + t_tri = bench(tri, x) + results["Triton"].append(t_tri) + print(f"{T:>6} {t_seq:>10.1f}ms {t_par:>8.1f}ms {t_tri:>6.1f}ms {t_seq/t_tri:>7.2f}x") + else: + print(f"{T:>6} {t_seq:>10.1f}ms {t_par:>8.1f}ms {'N/A':>8} {'N/A':>8}") + + # Plot + plt.figure(figsize=(8, 6)) + plt.plot(seq_lens, results["Sequential"], marker='o', color=WONG_PALETTE["vermillion"], label="Sequential (Baseline)") + plt.plot(seq_lens, results["Parallel"], marker='s', color=WONG_PALETTE["blue"], label="Parallel Scan (Mamba-like)") + + if HAS_TRITON and device == "cuda": + plt.plot(seq_lens, results["Triton"], marker='^', color=WONG_PALETTE["bluishgreen"], label="Triton Fused") + + plt.xlabel("Sequence Length (T)") + plt.ylabel("Inference Time (ms)") + plt.title("VLA Inference Time vs Sequence Length") + plt.legend() + plt.grid(True, linestyle='--', alpha=0.6) + plt.tight_layout() + plt.savefig("results/figures/speed_benchmark.png", dpi=300) + print("\nSaved speed benchmark plot to results/figures/speed_benchmark.png") + +if __name__ == "__main__": + main() diff --git a/scripts/plot_retrieval.py b/scripts/plot_retrieval.py index 96b8b91..908dc1f 100644 --- a/scripts/plot_retrieval.py +++ b/scripts/plot_retrieval.py @@ -4,59 +4,67 @@ import matplotlib.pyplot as plt import seaborn as sns +# Wong color palette (colorblind friendly) +WONG_PALETTE = { + "VLA": "#D55E00", # vermillion + "DeltaNet": "#0072B2", # blue + "Linear": "#E69F00", # orange + "VLAMamba": "#56B4E9", # skyblue + "VLATriton": "#009E73", # bluishgreen + "VLAKVExploding": "#CC79A7", # reddishpurple + "VLACombined": "#000000" # black +} + def main(): - # Find all CSV files generated by the retrieval benchmark - csv_files = glob.glob('results/retrieval_*.csv') - + csv_files = glob.glob("results/retrieval_*.csv") if not csv_files: - print("No CSV files found in 'results/'. Run the benchmark first.") + print("No retrieval CSV files found in results/.") return - - os.makedirs('plots', exist_ok=True) - sns.set_theme(style="darkgrid", palette="muted") - - # We will aggregate learning curves - # model -> seq_len -> dataframe - data_dict = {} - + + # Parse and combine + all_data = [] for f in csv_files: + # e.g. retrieval_VLA_256.csv basename = os.path.basename(f) - # Format: retrieval_{model}_{seq_len}.csv - parts = basename.replace('.csv', '').split('_') + parts = basename.replace(".csv", "").split("_") if len(parts) >= 3: model = parts[1] seq_len = int(parts[2]) + df = pd.read_csv(f) + df["Model"] = model + df["SeqLen"] = seq_len + all_data.append(df) - if model not in data_dict: - data_dict[model] = {} - data_dict[model][seq_len] = pd.read_csv(f) - - # Plot 1: Learning curves for a specific sequence length across all models - # Find all unique sequence lengths - all_seq_lens = set() - for m in data_dict: - all_seq_lens.update(data_dict[m].keys()) - - for seq_len in sorted(list(all_seq_lens)): - plt.figure(figsize=(10, 6)) - - for model in data_dict: - if seq_len in data_dict[model]: - df = data_dict[model][seq_len] - plt.plot(df['epoch'], df['test_acc'], label=f"{model}", marker='o', linewidth=2) - - plt.title(f'Associative Retrieval Test Accuracy (Sequence Length: {seq_len})', fontsize=14) - plt.xlabel('Epoch', fontsize=12) - plt.ylabel('Test Accuracy', fontsize=12) - plt.ylim(0, 1.05) - plt.legend(fontsize=12) - plt.grid(True, alpha=0.3) - plt.tight_layout() + if not all_data: + print("Could not parse any valid data.") + return + + combined_df = pd.concat(all_data, ignore_index=True) + + # Plot Accuracy + plt.figure(figsize=(10, 6)) + for model in combined_df["Model"].unique(): + model_data = combined_df[combined_df["Model"] == model] + # Calculate max accuracy per seq len to plot + acc_data = model_data.groupby("SeqLen")["test_acc"].max().reset_index() + acc_data = acc_data.sort_values(by="SeqLen") - plot_path = f'plots/retrieval_acc_seq_{seq_len}.png' - plt.savefig(plot_path, dpi=300) - print(f"Saved plot to {plot_path}") - plt.close() + color = WONG_PALETTE.get(model, "#000000") + plt.plot(acc_data["SeqLen"], acc_data["test_acc"], marker='o', label=model, color=color, linewidth=2) + + plt.xlabel("Sequence Length") + plt.ylabel("Test Accuracy (Exact Match)") + plt.title("Associative Retrieval Benchmark Accuracy") + plt.xscale('log', base=2) + plt.ylim(0, 1.05) + plt.grid(True, linestyle='--', alpha=0.6) + plt.legend() + plt.tight_layout() + + os.makedirs("results/benchmark_retrieval/images", exist_ok=True) + out_path = "results/benchmark_retrieval/images/retrieval_accuracy.png" + plt.savefig(out_path, dpi=300) + print(f"Saved accuracy plot to {out_path}") if __name__ == "__main__": main() diff --git a/src/models/attention/fast_vla.py b/src/models/attention/fast_vla.py new file mode 100644 index 0000000..70b3b5c --- /dev/null +++ b/src/models/attention/fast_vla.py @@ -0,0 +1,344 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from src.models.attention.penalty_builder import PenaltyBuilder + +# Try importing Triton +try: + import triton + import triton.language as tl + HAS_TRITON = True +except ImportError: + HAS_TRITON = False + +def _combine(F_l, G_l, F_r, G_r): + """Associative operator for the (F, G) recurrence.""" + return torch.matmul(F_r, F_l), torch.matmul(F_r, G_l) + G_r + +def parallel_scan(Fs, Gs): + """ + Inclusive parallel prefix scan for S_t = F_t @ S_{t-1} + G_t, S_0 = 0. + Args: + Fs, Gs : (T, B, d, d) + Returns: + S_all : (T, B, d, d) where S_all[t] = S_{t+1} + """ + T, B, d, _ = Fs.shape + + # Pad to next power of 2 + T_pad = 1 << (T - 1).bit_length() + pad = T_pad - T + if pad > 0: + I_p = torch.eye(d, device=Fs.device, dtype=Fs.dtype) \ + .unsqueeze(0).unsqueeze(0).expand(pad, B, -1, -1) + Z_p = torch.zeros(pad, B, d, d, device=Fs.device, dtype=Fs.dtype) + Fs = torch.cat([Fs, I_p], dim=0) + Gs = torch.cat([Gs, Z_p], dim=0) + + Fc, Gc = Fs.clone().to(Fs.dtype), Gs.clone().to(Gs.dtype) + + # ── Up-sweep ────────────────────────────────────────────────────────── + stride = 1 + while stride < T_pad: + r = torch.arange(2 * stride - 1, T_pad, 2 * stride, device=Fs.device) + l = r - stride + if r.numel() > 0: + Fc[r], Gc[r] = _combine(Fc[l], Gc[l], Fc[r], Gc[r]) + stride *= 2 + + # ── Down-sweep ──────────────────────────────────────────────────────── + Fc[T_pad - 1] = torch.eye(d, device=Fs.device, dtype=Fs.dtype) \ + .unsqueeze(0).expand(B, -1, -1) + Gc[T_pad - 1] = torch.zeros(B, d, d, device=Fs.device, dtype=Fs.dtype) + + stride = T_pad // 2 + while stride >= 1: + r = torch.arange(2 * stride - 1, T_pad, 2 * stride, device=Fs.device) + l = r - stride + if r.numel() > 0: + Fr, Gr = Fc[r].clone(), Gc[r].clone() + Fl, Gl = Fc[l].clone(), Gc[l].clone() + Fc[l], Gc[l] = Fr, Gr + Fc[r], Gc[r] = _combine(Fl, Gl, Fr, Gr) + stride //= 2 + + # Inclusive: apply original (Fs, Gs) to exclusive prefix + S_all = torch.matmul(Fs[:T], Gc[:T]) + Gs[:T] + return S_all + +def _compute_A_sequence(U_seq, lambda_0, d, stab_eps=1e-6, per_eps=1e-5, period=50): + """ + Sequential Sherman-Morrison: A_t for t = 1..T. + No .item() — fully torch.compile-able. + """ + B, T, _ = U_seq.shape + device, dtype = U_seq.device, U_seq.dtype + + I = torch.eye(d, device=device, dtype=dtype).unsqueeze(0).expand(B, -1, -1) + A_t = (1.0 / lambda_0) * I.clone() + A_all = [] + + for t in range(T): + u = U_seq[:, t, :] / math.sqrt(d) # (B, d) — normalise + uv = u.unsqueeze(-1) # (B, d, 1) + z = torch.bmm(A_t, uv) # (B, d, 1) + dot = torch.bmm(uv.transpose(1, 2), z).squeeze(-1).squeeze(-1) # (B,) + delt = torch.clamp(1.0 + dot, min=stab_eps) # (B,) + + upd = torch.bmm(z, z.transpose(1, 2)) / delt.view(B, 1, 1) # (B,d,d) + bad = (torch.abs(delt) < stab_eps).view(B, 1, 1) + A_t = torch.where(bad, A_t + stab_eps * I, A_t - upd) + + if (t + 1) % period == 0: + A_t = A_t + per_eps * I + + A_all.append(A_t.clone()) + + return torch.stack(A_all, dim=0) # (T, B, d, d) + + +if HAS_TRITON: + @triton.jit + def _sm_scan_kernel( + U_ptr, K_ptr, Out_ptr, + T, + stride_Ub, stride_Ut, + stride_Kb, stride_Kt, + stride_Ob, stride_Ot, + inv_lambda0: tl.constexpr, + inv_sqrt_D: tl.constexpr, + stab_eps: tl.constexpr, + period: tl.constexpr, + per_eps: tl.constexpr, + D: tl.constexpr, # power-of-2, <= 128 + ): + pid = tl.program_id(0) # one program per batch element + ids = tl.arange(0, D) # [0 .. D-1] + + # Initialise A = (1/lambda_0) * I in registers + A = tl.where( + ids[:, None] == ids[None, :], + inv_lambda0, 0.0, + ).to(tl.float32) + + # Pre-build stability / periodic nudge matrices + stab_mat = tl.where(ids[:, None] == ids[None, :], stab_eps, 0.0).to(tl.float32) + per_mat = tl.where(ids[:, None] == ids[None, :], per_eps, 0.0).to(tl.float32) + + U_base = U_ptr + pid * stride_Ub + K_base = K_ptr + pid * stride_Kb + O_base = Out_ptr + pid * stride_Ob + + for t in tl.range(T): + u = tl.load(U_base + t * stride_Ut + ids).to(tl.float32) * inv_sqrt_D + k = tl.load(K_base + t * stride_Kt + ids).to(tl.float32) + + z = tl.sum(A * u[None, :], axis=1) + + delta = tl.sum(u * z) + 1.0 + delta = tl.maximum(delta, stab_eps) + + outer = z[:, None] * z[None, :] + A_new = A - outer / delta + + A = tl.where(delta <= stab_eps, A + stab_mat, A_new) + + if (t + 1) % period == 0: + A = A + per_mat + + alpha = tl.sum(A * k[:, None], axis=0) + tl.store(O_base + t * stride_Ot + ids, alpha) + + def sm_scan_triton( + U: torch.Tensor, K: torch.Tensor, + lambda_0: float = 1.0, + stab_eps: float = 1e-6, + per_eps: float = 1e-5, + period: int = 50, + ) -> torch.Tensor: + B, T, D = U.shape + # Check power of 2 and <= 128 + assert D & (D - 1) == 0 and D <= 128, f"D must be power-of-2 <= 128, got {D}" + assert U.is_contiguous() and K.is_contiguous() + + alpha = torch.empty_like(U) + _sm_scan_kernel[(B,)]( + U, K, alpha, T, + U.stride(0), U.stride(1), + K.stride(0), K.stride(1), + alpha.stride(0), alpha.stride(1), + inv_lambda0 = 1.0 / lambda_0, + inv_sqrt_D = 1.0 / math.sqrt(D), + stab_eps = stab_eps, + period = period, + per_eps = per_eps, + D = D, + ) + return alpha + + +class _VLABase(nn.Module): + """Shared projections + S_t scan logic.""" + + def __init__(self, d_model, lambda_0=1.0, fixed_lambda=None, + stab_eps=1e-6, per_eps=1e-5, period=50, use_kv_exploding_fix=False): + super().__init__() + self.d_model = d_model + self.lambda_0 = lambda_0 + self.stab_eps = stab_eps + self.per_eps = per_eps + self.period = period + self.use_kv_exploding_fix = use_kv_exploding_fix + + self.W_q = nn.Linear(d_model, d_model) + self.W_k = nn.Linear(d_model, d_model) + self.W_v = nn.Linear(d_model, d_model) + self.out_norm = nn.LayerNorm(d_model) + self.W_o = nn.Linear(d_model, d_model) + self.penalty = PenaltyBuilder(d_model, fixed_lambda=fixed_lambda) + + def _build_S(self, alpha_all, K, V, d, B, T, device, dtype): + """Parallel scan for S_t given pre-computed alpha_all (T,B,d).""" + K_tr = K.permute(1, 0, 2) # (T, B, d) + V_tr = V.permute(1, 0, 2) + I_exp = torch.eye(d, device=device, dtype=dtype) \ + .unsqueeze(0).unsqueeze(0).expand(T, B, -1, -1) + + # Normalise k and alpha before rank-1 products to keep F_t a contraction. + # Without this, composing F matrices in the binary tree causes + # exponential blowup -> NaN gradients at long T. + # This acts as the KV Exploding fix. + if self.use_kv_exploding_fix: + k_n = F.normalize(K_tr, dim=-1) + a_n = F.normalize(alpha_all, dim=-1) + else: + k_n = K_tr + a_n = alpha_all + + k_alpha = torch.matmul(k_n.unsqueeze(-1), a_n.unsqueeze(-2)) # (T,B,d,d) + v_alpha = torch.matmul(V_tr.unsqueeze(-1), a_n.unsqueeze(-2)) # (T,B,d,d) + + # Clip G to prevent S explosion + vn = torch.norm(v_alpha, dim=(-2, -1), keepdim=True) + v_alpha = torch.where(vn > 10.0, v_alpha * 10.0 / (vn + 1e-6), v_alpha) + + Fs = I_exp - k_alpha + Gs = v_alpha + return parallel_scan(Fs, Gs) # (T, B, d, d) + + def _output(self, S_all, Q, d, B, T, dtype): + Q_tr = Q.permute(1, 0, 2).unsqueeze(-1) # (T, B, d, 1) + O = torch.matmul(S_all, Q_tr).squeeze(-1) \ + .permute(1, 0, 2).to(dtype) # (B, T, d) + return self.W_o(self.out_norm(O)) + + +class VLASequential(_VLABase): + """Baseline: Python for-loop over T (one CUDA kernel launch per token).""" + + def forward(self, x): + B, T, _ = x.shape + d = self.d_model + device, dtype = x.device, x.dtype + + Q = self.W_q(x) + K = self.W_k(x) + V = self.W_v(x) + + + _, U_seq, _ = self.penalty(K) + + I = torch.eye(d, device=device, dtype=dtype).unsqueeze(0).expand(B, -1, -1) + A_t = (1.0 / self.lambda_0) * I.clone() + S_t = torch.zeros(B, d, d, device=device, dtype=dtype) + outputs = [] + + for t in range(T): + u = U_seq[:, t, :].float() / math.sqrt(d) + uv = u.unsqueeze(-1) + z = torch.bmm(A_t, uv) + dot = torch.bmm(uv.transpose(1, 2), z).squeeze(-1).squeeze(-1) + delt = torch.clamp(1.0 + dot, min=self.stab_eps) + upd = torch.bmm(z, z.transpose(1, 2)) / delt.view(B, 1, 1) + bad = (torch.abs(delt) < self.stab_eps).view(B, 1, 1) + A_t = torch.where(bad, A_t + self.stab_eps * I, A_t - upd) + if (t + 1) % self.period == 0: + A_t = A_t + self.per_eps * I + + k_t = K[:, t, :].float() + v_t = V[:, t, :].float() + q_t = Q[:, t, :].float() + alpha_t = torch.bmm(A_t, k_t.unsqueeze(-1)).squeeze(-1) + + v_hat = torch.bmm(S_t, k_t.unsqueeze(-1)).squeeze(-1) + e_t = v_t - v_hat + en = torch.norm(e_t, dim=-1, keepdim=True) + e_t = torch.where(en > 10.0, e_t * 10.0 / (en + 1e-6), e_t) + + S_t = S_t + torch.matmul(e_t.unsqueeze(2), alpha_t.unsqueeze(1)) + outputs.append(torch.matmul(S_t, q_t.unsqueeze(2)).squeeze(2)) + + O = torch.stack(outputs, dim=1).to(dtype) + return self.W_o(self.out_norm(O)) + + +class VLAParallel(_VLABase): + """Python A_t loop + parallel prefix scan for S_t.""" + + def forward(self, x): + B, T, _ = x.shape + d = self.d_model + device, dtype = x.device, x.dtype + + Q = self.W_q(x) + K = self.W_k(x) + V = self.W_v(x) + + + _, U_seq, _ = self.penalty(K) + + A_all = _compute_A_sequence(U_seq.float(), self.lambda_0, d, + self.stab_eps, self.per_eps, self.period) + K_t = K.permute(1, 0, 2).unsqueeze(-1) + # Use a_t_scaled computation from VLA base if needed, but here alpha_t = A_t @ k_t + alpha_all = torch.matmul(A_all, K_t).squeeze(-1) # (T, B, d) + + S_all = self._build_S(alpha_all, K, V, d, B, T, device, dtype) + return self._output(S_all, Q, d, B, T, dtype) + + +class VLATriton(_VLABase): + """ + Triton kernel fuses entire A_t loop into ONE kernel launch. + """ + + def forward(self, x): + assert HAS_TRITON and x.device.type == "cuda", \ + "VLATriton requires CUDA + Triton" + B, T, _ = x.shape + d = self.d_model + device, dtype = x.device, x.dtype + + Q = self.W_q(x) + K = self.W_k(x) + V = self.W_v(x) + + + _, U_seq, _ = self.penalty(K) + + U_c = U_seq.float().contiguous() + K_c = K.float().contiguous() + alpha_all = sm_scan_triton( + U_c, K_c, + lambda_0 = self.lambda_0, + stab_eps = self.stab_eps, + per_eps = self.per_eps, + period = self.period, + ) # (B, T, d) + alpha_all = alpha_all.permute(1, 0, 2) # -> (T, B, d) + + S_all = self._build_S(alpha_all, K, V, d, B, T, device, dtype) + return self._output(S_all, Q, d, B, T, dtype) + + diff --git a/tests/test_fast_vla.py b/tests/test_fast_vla.py new file mode 100644 index 0000000..3bdc87f --- /dev/null +++ b/tests/test_fast_vla.py @@ -0,0 +1,22 @@ +import torch +import pytest +from src.models.attention.fast_vla import VLAParallel, VLATriton, HAS_TRITON + +def test_fast_vla_variants(): + d = 64 + B = 2 + T = 16 + device = "cpu" + + x = torch.randn(B, T, d) + + par1 = VLAParallel(d_model=d, use_kv_exploding_fix=True) + out1 = par1(x) + assert out1.shape == (B, T, d) + + par2 = VLAParallel(d_model=d, use_kv_exploding_fix=False) + out2 = par2(x) + assert out2.shape == (B, T, d) + +if __name__ == "__main__": + pytest.main([__file__])