test: probe torch.matmul backend + ixformer.matmul/linear + Python loop overhead

qwen3_6_scripts/probe_torch_matmul_backend.sh

@@ -0,0 +1,172 @@
+#!/bin/bash
+# probe_torch_matmul_backend.sh — Find out what torch.matmul actually calls on BI-V100
+
+echo "=== 1. torch.matmul profiling — which kernel launches ==="
+python3 << 'PY'
+import torch
+import os
+
+M, N, K = 256, 11008, 4096
+A = torch.randn(M, K, dtype=torch.float16, device='cuda')
+B = torch.randn(K, N, dtype=torch.float16, device='cuda')
+
+# Warmup
+for _ in range(5):
+    torch.matmul(A, B)
+torch.cuda.synchronize()
+
+# Profile with torch profiler
+with torch.profiler.profile(
+    activities=[torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.CUDA],
+    record_shapes=True,
+) as prof:
+    for _ in range(3):
+        torch.matmul(A, B)
+    torch.cuda.synchronize()
+
+print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
+PY
+
+echo ""
+echo "=== 2. CUDA_LAUNCH_BLOCKING trace ==="
+CUDA_LAUNCH_BLOCKING=1 python3 << 'PY'
+import torch
+import os
+
+M, N, K = 256, 11008, 4096
+A = torch.randn(M, K, dtype=torch.float16, device='cuda')
+B = torch.randn(K, N, dtype=torch.float16, device='cuda')
+
+# Warmup
+torch.matmul(A, B)
+torch.cuda.synchronize()
+
+# Single call with sync
+torch.cuda.synchronize()
+C = torch.matmul(A, B)
+torch.cuda.synchronize()
+print("torch.matmul completed")
+print(f"Result shape: {C.shape}, dtype: {C.dtype}")
+PY
+
+echo ""
+echo "=== 3. Check if ixformer.matmul is faster or same as torch.matmul ==="
+python3 << 'PY'
+import torch
+import time
+
+M, N, K = 256, 11008, 4096
+A = torch.randn(M, K, dtype=torch.float16, device='cuda')
+B = torch.randn(K, N, dtype=torch.float16, device='cuda')
+
+# torch.matmul
+for _ in range(10):
+    torch.matmul(A, B)
+torch.cuda.synchronize()
+t0 = time.time()
+for _ in range(100):
+    torch.matmul(A, B)
+torch.cuda.synchronize()
+ms_torch = (time.time() - t0) / 100 * 1000
+
+# torch.mm (should be same)
+for _ in range(10):
+    torch.mm(A, B)
+torch.cuda.synchronize()
+t0 = time.time()
+for _ in range(100):
+    torch.mm(A, B)
+torch.cuda.synchronize()
+ms_mm = (time.time() - t0) / 100 * 1000
+
+# ixformer.matmul
+import ixformer
+A_ix = ixformer.Tensor(A)
+B_ix = ixformer.Tensor(B)
+for _ in range(10):
+    ixformer.matmul(A_ix, B_ix)
+torch.cuda.synchronize()
+t0 = time.time()
+for _ in range(100):
+    ixformer.matmul(A_ix, B_ix)
+torch.cuda.synchronize()
+ms_ix = (time.time() - t0) / 100 * 1000
+
+# ixformer.linear
+A2 = torch.randn(M, K, dtype=torch.float16, device='cuda')
+W = torch.randn(N, K, dtype=torch.float16, device='cuda')  # (out, in)
+A2_ix = ixformer.Tensor(A2)
+W_ix = ixformer.Tensor(W)
+for _ in range(10):
+    ixformer.linear(A2_ix, W_ix)
+torch.cuda.synchronize()
+t0 = time.time()
+for _ in range(100):
+    ixformer.linear(A2_ix, W_ix)
+torch.cuda.synchronize()
+ms_linear = (time.time() - t0) / 100 * 1000
+
+# F.linear (torch)
+for _ in range(10):
+    torch.nn.functional.linear(A2, W)
+torch.cuda.synchronize()
+t0 = time.time()
+for _ in range(100):
+    torch.nn.functional.linear(A2, W)
+torch.cuda.synchronize()
+ms_flinear = (time.time() - t0) / 100 * 1000
+
+# cublas via direct cublasHgemm (through torch C++ extension would need compile)
+# Skip for now
+
+print(f"torch.matmul:       {ms_torch:.3f} ms")
+print(f"torch.mm:           {ms_mm:.3f} ms")
+print(f"ixformer.matmul:    {ms_ix:.3f} ms")
+print(f"ixformer.linear:    {ms_linear:.3f} ms")
+print(f"F.linear:           {ms_flinear:.3f} ms")
+PY
+
+echo ""
+echo "=== 4. MoE hot loop: Python for-loop overhead measurement ==="
+python3 << 'PY'
+import torch
+import time
+
+num_experts = 8
+top_k = 8
+K, N = 4096, 11008
+
+# Simulate decode: 1 token, 8 experts
+hidden = torch.randn(1, K, dtype=torch.float16, device='cuda')
+weights = [torch.randn(N, K, dtype=torch.float16, device='cuda') for _ in range(num_experts)]
+
+# Warmup
+for w in weights:
+    torch.nn.functional.linear(hidden, w)
+torch.cuda.synchronize()
+
+# Measure: 8 F.linear calls in Python loop
+t0 = time.time()
+for _ in range(1000):
+    results = []
+    for i in range(top_k):
+        results.append(torch.nn.functional.linear(hidden, weights[i]))
+    out = sum(results)
+torch.cuda.synchronize()
+ms_loop = (time.time() - t0) / 1000 * 1000
+
+# Measure: single F.linear with same total FLOPS
+big_w = torch.randn(N * top_k, K, dtype=torch.float16, device='cuda')
+big_hidden = hidden.expand(top_k, K).contiguous().view(top_k, K)
+t0 = time.time()
+for _ in range(1000):
+    out = torch.mm(big_hidden, big_w.t())
+torch.cuda.synchronize()
+ms_single = (time.time() - t0) / 1000 * 1000
+
+print(f"8x F.linear loop (decode): {ms_loop:.3f} ms")
+print(f"1x torch.mm equivalent:    {ms_single:.3f} ms")
+print(f"Python loop overhead:      {ms_loop - ms_single:.3f} ms")
+print(f"Per-expert overhead:       {(ms_loop - ms_single)/top_k:.3f} ms")
+PY