Upload folder using huggingface_hub

README.md

@@ -93,7 +93,8 @@ B = 2048
 K = 256
 F = 1024 * 128
 D = 1024
-warmup = 5
+WARMUP = 5
+NUM_ITER = 100
 dtype = torch.float32
 
 vals = None
@@ -126,7 +127,7 @@ def zero_grad():
     torch.cuda.empty_cache()
 
 
-for i in range(100 + warmup):
+for i in range(NUM_ITER + WARMUP):
     init_parameters()
     start_kernel = torch.cuda.Event(enable_timing=True)
     end_kernel = torch.cuda.Event(enable_timing=True)
@@ -143,7 +144,7 @@ for i in range(100 + warmup):
     loss_vanilla = hierarchical_sae_loss(indices, decoder, vals, bias, target)
     loss_vanilla.backward()
     end_vanilla.record()
-    if i >= warmup:
+    if i >= WARMUP:
         torch.cuda.synchronize()
         timing_kernel.append(start_kernel.elapsed_time(end_kernel))
         timing_vanilla.append(start_vanilla.elapsed_time(end_vanilla))

example.py

@@ -31,7 +31,8 @@ B = 2048
 K = 256
 F = 1024 * 128
 D = 1024
-warmup = 5
+WARMUP = 5
+NUM_ITER = 100
 dtype = torch.float32
 
 vals = None
@@ -64,7 +65,7 @@ def zero_grad():
     torch.cuda.empty_cache()
 
 
-for i in range(100 + warmup):
+for i in range(NUM_ITER + WARMUP):
     init_parameters()
     start_kernel = torch.cuda.Event(enable_timing=True)
     end_kernel = torch.cuda.Event(enable_timing=True)
@@ -81,7 +82,7 @@ for i in range(100 + warmup):
     loss_vanilla = hierarchical_sae_loss(indices, decoder, vals, bias, target)
     loss_vanilla.backward()
     end_vanilla.record()
-    if i >= warmup:
+    if i >= WARMUP:
         torch.cuda.synchronize()
         timing_kernel.append(start_kernel.elapsed_time(end_kernel))
         timing_vanilla.append(start_vanilla.elapsed_time(end_vanilla))
@@ -97,4 +98,4 @@ else:
 
 print(f"🦎 Triton Kernel Time (Ours): {np.mean(timing_kernel):.4f} ± {np.std(timing_kernel):.4f} ms")
 print(f"🔥 Torch Compile Kernel Time: {np.mean(timing_vanilla):.4f} ± {np.std(timing_vanilla):.4f} ms")
-print(f"🚀 Speedup: {np.mean(timing_vanilla) / np.mean(timing_kernel):.2f}x")
+print(f"🚀 Speedup: {np.mean(timing_vanilla) / np.mean(timing_kernel):.2f}x")