Upload folder using huggingface_hub

README.md

@@ -4,4 +4,6 @@ tags:
 - kernel
 ---
 
-# batch_invariant_kernel
+# batch_invariant_kernel
+
+To try out the example of the code 

readme_example.py

@@ -0,0 +1,202 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "torch",
+#     "numpy",
+#     "kernels",
+# ]
+# ///
+
+import torch
+from kernels import get_kernel
+
+# Load batch_invariant_kernel via kernels library
+batch_invariant_kernel = get_kernel("gagan3012/batch_invariant_kernel")
+
+# Set device and seed for reproducibility
+device = "cuda"
+torch.manual_seed(42)
+torch.cuda.manual_seed(42)
+
+print("🚀 Testing batch_invariant_kernel from Hugging Face Hub")
+print(f"✅ CUDA is available. Using device: {torch.cuda.get_device_name()}")
+
+# Test 1: Matrix Multiplication
+print("\n" + "=" * 60)
+print("🧪 Test 1: Persistent Matrix Multiplication")
+print("=" * 60)
+
+# Parameters for matrix multiplication
+M, K, N = 512, 256, 1024
+a = torch.randn(M, K, device=device, dtype=torch.float32)
+b = torch.randn(K, N, device=device, dtype=torch.float32)
+bias = torch.randn(N, device=device, dtype=torch.float32)
+
+print(f"Matrix A shape: {a.shape}")
+print(f"Matrix B shape: {b.shape}")
+print(f"Bias shape: {bias.shape}")
+
+# Run matrix multiplication without bias
+start_event = torch.cuda.Event(enable_timing=True)
+end_event = torch.cuda.Event(enable_timing=True)
+
+start_event.record()
+output_no_bias = batch_invariant_kernel.matmul_persistent(a, b)
+end_event.record()
+torch.cuda.synchronize()
+time_no_bias = start_event.elapsed_time(end_event)
+
+print(f"\nMatrix multiplication (no bias) completed!")
+print(f"Output shape: {output_no_bias.shape}")
+print(f"Execution time: {time_no_bias:.3f} ms")
+
+# Run matrix multiplication with bias
+start_event.record()
+output_with_bias = batch_invariant_kernel.matmul_persistent(a, b, bias)
+end_event.record()
+torch.cuda.synchronize()
+time_with_bias = start_event.elapsed_time(end_event)
+
+print(f"\nMatrix multiplication (with bias) completed!")
+print(f"Output shape: {output_with_bias.shape}")
+print(f"Execution time: {time_with_bias:.3f} ms")
+
+# Verify correctness
+expected_no_bias = torch.mm(a, b)
+expected_with_bias = torch.mm(a, b) + bias
+
+max_diff_no_bias = torch.max(torch.abs(output_no_bias - expected_no_bias)).item()
+max_diff_with_bias = torch.max(torch.abs(output_with_bias - expected_with_bias)).item()
+
+print(f"Max difference (no bias): {max_diff_no_bias:.6f}")
+print(f"Max difference (with bias): {max_diff_with_bias:.6f}")
+
+# Test 2: Log Softmax
+print("\n" + "=" * 60)
+print("🧪 Test 2: Log Softmax")
+print("=" * 60)
+
+# Parameters for log softmax (typical attention dimensions)
+batch_size = 4
+seq_len = 512
+vocab_size = 32000
+
+logits = torch.randn(
+    batch_size, seq_len, vocab_size, device=device, dtype=torch.float32
+)
+print(f"Input logits shape: {logits.shape}")
+
+# Run log softmax
+start_event.record()
+log_probs = batch_invariant_kernel.log_softmax(logits, dim=-1)
+end_event.record()
+torch.cuda.synchronize()
+time_log_softmax = start_event.elapsed_time(end_event)
+
+print(f"\nLog softmax completed!")
+print(f"Output shape: {log_probs.shape}")
+print(f"Execution time: {time_log_softmax:.3f} ms")
+
+# Verify correctness
+expected_log_probs = torch.log_softmax(logits, dim=-1)
+max_diff_log_softmax = torch.max(torch.abs(log_probs - expected_log_probs)).item()
+print(f"Max difference vs PyTorch: {max_diff_log_softmax:.6f}")
+
+# Test 3: Mean Reduction
+print("\n" + "=" * 60)
+print("🧪 Test 3: Mean Dimension Reduction")
+print("=" * 60)
+
+# Parameters for mean reduction (typical layer norm dimensions)
+batch_size = 8
+seq_len = 256
+hidden_size = 768
+
+hidden_states = torch.randn(
+    batch_size, seq_len, hidden_size, device=device, dtype=torch.float32
+)
+print(f"Input hidden states shape: {hidden_states.shape}")
+
+# Test reduction along different dimensions
+for dim in [0, 1, 2]:
+    start_event.record()
+    mean_output = batch_invariant_kernel.mean_dim(hidden_states, dim=dim, keepdim=False)
+    end_event.record()
+    torch.cuda.synchronize()
+    time_mean = start_event.elapsed_time(end_event)
+
+    expected_mean = torch.mean(hidden_states, dim=dim, keepdim=False)
+    max_diff_mean = torch.max(torch.abs(mean_output - expected_mean)).item()
+
+    print(f"\nMean reduction along dim {dim}:")
+    print(f"  Output shape: {mean_output.shape}")
+    print(f"  Execution time: {time_mean:.3f} ms")
+    print(f"  Max difference vs PyTorch: {max_diff_mean:.6f}")
+
+# Test 4: End-to-End Attention-like Computation
+print("\n" + "=" * 60)
+print("🧪 Test 4: End-to-End Attention-like Computation")
+print("=" * 60)
+
+# Simulate a simple attention computation using our kernels
+batch_size = 4
+seq_len = 128
+hidden_size = 512
+num_heads = 8
+head_dim = hidden_size // num_heads
+
+# Input embeddings
+x = torch.randn(batch_size, seq_len, hidden_size, device=device, dtype=torch.float32)
+
+# Weight matrices for Q, K, V projections
+w_q = torch.randn(hidden_size, hidden_size, device=device, dtype=torch.float32)
+w_k = torch.randn(hidden_size, hidden_size, device=device, dtype=torch.float32)
+w_v = torch.randn(hidden_size, hidden_size, device=device, dtype=torch.float32)
+w_o = torch.randn(hidden_size, hidden_size, device=device, dtype=torch.float32)
+
+print(f"Input shape: {x.shape}")
+print("Computing Q, K, V projections using batch_invariant matmul...")
+
+# Reshape for batch matrix multiplication
+x_flat = x.view(-1, hidden_size)  # (batch_size * seq_len, hidden_size)
+
+start_event.record()
+
+# Compute Q, K, V using our custom matmul
+q_flat = batch_invariant_kernel.matmul_persistent(x_flat, w_q)
+k_flat = batch_invariant_kernel.matmul_persistent(x_flat, w_k)
+v_flat = batch_invariant_kernel.matmul_persistent(x_flat, w_v)
+
+# Reshape to multi-head format
+q = q_flat.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
+k = k_flat.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
+v = v_flat.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
+
+# Compute attention scores
+scores = torch.matmul(q, k.transpose(-2, -1)) / (head_dim**0.5)
+
+# Apply softmax using our custom log_softmax (convert to softmax)
+log_attn_weights = batch_invariant_kernel.log_softmax(scores, dim=-1)
+attn_weights = torch.exp(log_attn_weights)
+
+# Apply attention to values
+attn_output = torch.matmul(attn_weights, v)
+
+# Reshape and apply output projection
+attn_output = (
+    attn_output.transpose(1, 2).contiguous().view(batch_size * seq_len, hidden_size)
+)
+final_output = batch_invariant_kernel.matmul_persistent(attn_output, w_o)
+final_output = final_output.view(batch_size, seq_len, hidden_size)
+
+end_event.record()
+torch.cuda.synchronize()
+total_time = start_event.elapsed_time(end_event)
+
+print(f"\nEnd-to-end attention computation completed!")
+print(f"Final output shape: {final_output.shape}")
+print(f"Total execution time: {total_time:.3f} ms")
+print(
+    f"Output tensor stats - Mean: {final_output.mean().item():.4f}, Std: {final_output.std().item():.4f}"
+)
+