Skip to content

Implement fast VLA attention with Triton and update benchmarks #10

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Vishal-sys-code merged 2 commits into from
Apr 12, 2026
Merged

Implement fast VLA attention with Triton and update benchmarks #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0b0ef8f
Update benchmark_retrieval.py
Vishal-sys-code Apr 12, 2026
79f5ce0
Implement fast VLA attention with Triton
Vishal-sys-code Apr 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions scripts/benchmark_retrieval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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")
Expand All @@ -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:
Expand Down
88 changes: 88 additions & 0 deletions scripts/benchmark_speed.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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()
94 changes: 51 additions & 43 deletions scripts/plot_retrieval.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Loading
Loading