Source snapshot of github.com/Scottcjn/vulkan-matmul@036b448

35d65e4 verified 13 days ago

9.47 kB

#!/usr/bin/env python3
"""
test_client.py v2 — Tests VRAM weight caching in vulkan_matmul_server
Protocol v2 (32-byte header):
  magic(4) M(4) N(4) K(4) type(4) a_only(4) key_lo(4) key_hi(4)
  type=0 plain matmul (A+B → C)
  type=1 cached matmul
    a_only=0: send A+B (weight will be cached on server)
    a_only=1: send A only (weight should already be in VRAM cache)
Response: magic(4) status(4) M(4) N(4) then C data
  status=0  OK
  status=2  NEED_WEIGHT (evicted, must retry with a_only=0)
"""

import socket, struct, numpy as np, time, sys
from collections import defaultdict

MAGIC         = 0x564B4D54
STATUS_OK     = 0
STATUS_NEED_W = 2
TYPE_PLAIN    = 0
TYPE_CACHED   = 1

HOST = "127.0.0.1"
PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 8097

# ─── Client-side cache tracking ───────────────────────────────
# Maps weight_key → (K, N) for weights we believe are cached on server
_warm_keys = {}

# Monotonic key counter — avoids id(B) address-reuse collisions.
# Python may reuse the same memory address for different numpy arrays if the
# old array is freed before the new one is allocated (same shape/size).
# Using id(B) as a cache key then causes the server to see a "cache hit" for
# a completely different weight matrix → protocol desync ("bad magic").
_key_seq = 0
def _next_key():
    global _key_seq
    _key_seq += 1
    return _key_seq

def _connect():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect((HOST, PORT))
    s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    return s

# Persistent connection for this test session
_sock = _connect()

def _recv_all(n):
    buf = b''
    while len(buf) < n:
        chunk = _sock.recv(min(65536, n - len(buf)))
        if not chunk:
            raise ConnectionError("Server closed")
        buf += chunk
    return buf

def _send_all(data):
    mv = memoryview(data)
    sent = 0
    while sent < len(mv):
        n = _sock.send(mv[sent:])
        if n == 0:
            raise ConnectionError("Send failed")
        sent += n

def vk_matmul_cached(M, N, K, A, B, weight_key):
    """
    Cached matmul: C[M,N] = A[M,K] × B[K,N]
    On cache hit: only sends A (28KB for M=1 gen)
    On cache miss: sends A+B, server caches B in VRAM
    """
    global _warm_keys

    a_only = 1 if weight_key in _warm_keys else 0

    key_lo = weight_key & 0xFFFFFFFF
    key_hi = (weight_key >> 32) & 0xFFFFFFFF

    hdr = struct.pack('<IIIIIIII', MAGIC, M, N, K, TYPE_CACHED, a_only, key_lo, key_hi)
    _send_all(hdr)
    _send_all(A.astype(np.float32).tobytes())
    if not a_only:
        _send_all(B.astype(np.float32).tobytes())

    # Response
    rhdr = _recv_all(16)
    magic, status, rM, rN = struct.unpack('<IIII', rhdr)
    assert magic == MAGIC

    if status == STATUS_NEED_W:
        # Server evicted our key — retry with full B
        del _warm_keys[weight_key]
        return vk_matmul_cached(M, N, K, A, B, weight_key)

    assert status == STATUS_OK, f"status={status}"

    data = _recv_all(M * N * 4)
    C = np.frombuffer(data, dtype=np.float32).reshape(M, N)

    # Mark weight as cached on server
    if not a_only:
        _warm_keys[weight_key] = (K, N)

    return C

def vk_matmul_plain(M, N, K, A, B):
    """Plain (uncached) matmul for comparison."""
    hdr = struct.pack('<IIIIIIII', MAGIC, M, N, K, TYPE_PLAIN, 0, 0, 0)
    _send_all(hdr)
    _send_all(A.astype(np.float32).tobytes())
    _send_all(B.astype(np.float32).tobytes())
    rhdr = _recv_all(16)
    magic, status, rM, rN = struct.unpack('<IIII', rhdr)
    assert magic == MAGIC and status == 0
    return np.frombuffer(_recv_all(M*N*4), dtype=np.float32).reshape(M, N)

# ─── Tests ────────────────────────────────────────────────────

print(f"=== Vulkan Matmul Server v2 Test — VRAM Cache ===")
print(f"Connecting to {HOST}:{PORT}...\n")

print("── 1. Correctness (plain vs cached) ──")
for M, N, K, label in [
    (1,    256,  256,  "tiny"),
    (1,    1024, 1024, "1K gen"),
    (128,  2048, 2048, "128 prefill"),
    (1,    2048, 7168, "MoE expert"),
]:
    A = np.random.randn(M, K).astype(np.float32)
    B = np.random.randn(K, N).astype(np.float32)
    key = _next_key()

    C_np = A @ B
    C_vk = vk_matmul_cached(M, N, K, A, B, key)
    err  = np.max(np.abs(C_vk - C_np))
    print(f"  {'✓' if err < 0.01 else '✗'} [{M}×{N}×{K}] {label:20s}  maxerr={err:.2e}")

print()
print("── 2. Cache speedup — simulating 120B generation tokens ──")
print("   (Upload weight ONCE, then send only activations for each token)")
print()

shapes = [
    (7168, 7168,  "q_proj/k_proj/v_proj  (attention)"),
    (7168, 7168,  "o_proj                (attention)"),
    (7168, 2048,  "expert gate_proj      (MoE FFN)"),
    (7168, 2048,  "expert up_proj        (MoE FFN)"),
    (2048, 7168,  "expert down_proj      (MoE FFN)"),
]
N_TOKENS = 20

for (K, N, label) in shapes:
    B = np.random.randn(K, N).astype(np.float32)
    key = _next_key()   # stable unique key — no id() address-reuse collisions
    # Remove from warm cache so first call uploads (key is always fresh here)
    _warm_keys.pop(key, None)

    cold_times = []
    warm_times = []

    for tok in range(N_TOKENS):
        A = np.random.randn(1, K).astype(np.float32)  # single token
        t0 = time.time()
        C = vk_matmul_cached(1, N, K, A, B, key)
        elapsed = (time.time() - t0) * 1000
        if tok == 0:
            cold_times.append(elapsed)  # first call = upload + compute
        else:
            warm_times.append(elapsed)  # subsequent = activation only + compute

    cold_avg = sum(cold_times)
    warm_avg = sum(warm_times) / len(warm_times)
    speedup  = cold_avg / warm_avg if warm_avg > 0 else 0
    wt_mb    = K * N * 4 / 1e6

    print(f"  [{K}×{N}]  {label}")
    print(f"    Cold (upload {wt_mb:.0f}MB weight): {cold_avg:.0f}ms")
    print(f"    Warm (send 4KB activation only):   {warm_avg:.1f}ms  ← {speedup:.0f}× faster")
    print()

print()
print("── 3. Throughput: tokens/sec with cached weights ──")
# Simulate a full 120B MoE forward pass for one token
# Real 120B MoE: ~80 layers, each with 4 attn projections + 2 active experts × 3 projections
LAYERS      = 4   # subset for timing
ATTN_PROJS  = 4   # Q,K,V,O
EXPERTS     = 2   # active experts per token
EXP_PROJS   = 3   # gate, up, down

# Pre-create weight matrices (simulating model loaded in POWER8 RAM)
attn_weights = [np.random.randn(7168, 7168).astype(np.float32) for _ in range(ATTN_PROJS)]
exp_weights  = [(np.random.randn(7168, 2048).astype(np.float32),
                 np.random.randn(7168, 2048).astype(np.float32),
                 np.random.randn(2048, 7168).astype(np.float32))
                for _ in range(EXPERTS)]

# Warm up the cache (first forward pass)
print("  Warming up weight cache (first token upload)...")
t_warmup = time.time()
A_dummy = np.random.randn(1, 7168).astype(np.float32)
for layer in range(LAYERS):
    for wi, W in enumerate(attn_weights):
        key = (layer * 1000 + wi) & 0xFFFFFFFFFFFFFFFF
        _warm_keys.pop(key, None)
        vk_matmul_cached(1, 7168, 7168, A_dummy, W, key)
    for ei, (Wg, Wu, Wd) in enumerate(exp_weights):
        kg = (layer * 1000 + 100 + ei * 3) & 0xFFFFFFFFFFFFFFFF
        ku = (layer * 1000 + 101 + ei * 3) & 0xFFFFFFFFFFFFFFFF
        kd = (layer * 1000 + 102 + ei * 3) & 0xFFFFFFFFFFFFFFFF
        for k, W, K, N in [(kg,Wg,7168,2048),(ku,Wu,7168,2048),(kd,Wd,2048,7168)]:
            _warm_keys.pop(k, None)
            vk_matmul_cached(1, N, K, np.random.randn(1,K).astype(np.float32), W, k)
warmup_ms = (time.time() - t_warmup) * 1000
print(f"  Warmup done in {warmup_ms:.0f}ms\n")

# Now time N tokens with warm cache
N_TIME_TOKENS = 10
t_start = time.time()
for tok in range(N_TIME_TOKENS):
    A = np.random.randn(1, 7168).astype(np.float32)
    for layer in range(LAYERS):
        for wi, W in enumerate(attn_weights):
            key = (layer * 1000 + wi) & 0xFFFFFFFFFFFFFFFF
            vk_matmul_cached(1, 7168, 7168, A, W, key)
        for ei, (Wg, Wu, Wd) in enumerate(exp_weights):
            kg = (layer * 1000 + 100 + ei * 3) & 0xFFFFFFFFFFFFFFFF
            ku = (layer * 1000 + 101 + ei * 3) & 0xFFFFFFFFFFFFFFFF
            kd = (layer * 1000 + 102 + ei * 3) & 0xFFFFFFFFFFFFFFFF
            vk_matmul_cached(1, 2048, 7168, np.random.randn(1,7168).astype(np.float32), Wg, kg)
            vk_matmul_cached(1, 2048, 7168, np.random.randn(1,7168).astype(np.float32), Wu, ku)
            vk_matmul_cached(1, 7168, 2048, np.random.randn(1,2048).astype(np.float32), Wd, kd)

elapsed_s = time.time() - t_start
ms_per_tok = elapsed_s * 1000 / N_TIME_TOKENS
ops_per_tok = LAYERS * (ATTN_PROJS + EXPERTS * EXP_PROJS)

print(f"  {N_TIME_TOKENS} tokens × {LAYERS} layers × "
      f"({ATTN_PROJS} attn + {EXPERTS}×{EXP_PROJS} MoE) ops")
print(f"  Time: {elapsed_s*1000:.0f}ms total  =  {ms_per_tok:.1f}ms/token")
print(f"  Throughput: {1000/ms_per_tok:.2f} tokens/sec  "
      f"({ops_per_tok} matmuls/token, {ms_per_tok/ops_per_tok:.1f}ms each)")
print(f"\n  Cache stats: {len(_warm_keys)} weights resident in VRAM")
print(f"  (Extrapolated full 80-layer 120B: "
      f"~{ms_per_tok * 80/LAYERS:.0f}ms/token = "
      f"~{1000/(ms_per_tok * 80/LAYERS):.2f} t/s)")

_sock.close()
print("\nAll done.")

	#!/usr/bin/env python3
	"""
	test_client.py v2 — Tests VRAM weight caching in vulkan_matmul_server

	Protocol v2 (32-byte header):
	magic(4) M(4) N(4) K(4) type(4) a_only(4) key_lo(4) key_hi(4)
	type=0 plain matmul (A+B → C)
	type=1 cached matmul
	a_only=0: send A+B (weight will be cached on server)
	a_only=1: send A only (weight should already be in VRAM cache)

	Response: magic(4) status(4) M(4) N(4) then C data
	status=0 OK
	status=2 NEED_WEIGHT (evicted, must retry with a_only=0)
	"""

	import socket, struct, numpy as np, time, sys
	from collections import defaultdict

	MAGIC = 0x564B4D54
	STATUS_OK = 0
	STATUS_NEED_W = 2
	TYPE_PLAIN = 0
	TYPE_CACHED = 1

	HOST = "127.0.0.1"
	PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 8097

	# ─── Client-side cache tracking ───────────────────────────────
	# Maps weight_key → (K, N) for weights we believe are cached on server
	_warm_keys = {}

	# Monotonic key counter — avoids id(B) address-reuse collisions.
	# Python may reuse the same memory address for different numpy arrays if the
	# old array is freed before the new one is allocated (same shape/size).
	# Using id(B) as a cache key then causes the server to see a "cache hit" for
	# a completely different weight matrix → protocol desync ("bad magic").
	_key_seq = 0
	def _next_key():
	global _key_seq
	_key_seq += 1
	return _key_seq

	def _connect():
	s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	s.connect((HOST, PORT))
	s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
	return s

	# Persistent connection for this test session
	_sock = _connect()

	def _recv_all(n):
	buf = b''
	while len(buf) < n:
	chunk = _sock.recv(min(65536, n - len(buf)))
	if not chunk:
	raise ConnectionError("Server closed")
	buf += chunk
	return buf

	def _send_all(data):
	mv = memoryview(data)
	sent = 0
	while sent < len(mv):
	n = _sock.send(mv[sent:])
	if n == 0:
	raise ConnectionError("Send failed")
	sent += n

	def vk_matmul_cached(M, N, K, A, B, weight_key):
	"""
	Cached matmul: C[M,N] = A[M,K] × B[K,N]
	On cache hit: only sends A (28KB for M=1 gen)
	On cache miss: sends A+B, server caches B in VRAM
	"""
	global _warm_keys

	a_only = 1 if weight_key in _warm_keys else 0

	key_lo = weight_key & 0xFFFFFFFF
	key_hi = (weight_key >> 32) & 0xFFFFFFFF

	hdr = struct.pack('<IIIIIIII', MAGIC, M, N, K, TYPE_CACHED, a_only, key_lo, key_hi)
	_send_all(hdr)
	_send_all(A.astype(np.float32).tobytes())
	if not a_only:
	_send_all(B.astype(np.float32).tobytes())

	# Response
	rhdr = _recv_all(16)
	magic, status, rM, rN = struct.unpack('<IIII', rhdr)
	assert magic == MAGIC

	if status == STATUS_NEED_W:
	# Server evicted our key — retry with full B
	del _warm_keys[weight_key]
	return vk_matmul_cached(M, N, K, A, B, weight_key)

	assert status == STATUS_OK, f"status={status}"

	data = _recv_all(M * N * 4)
	C = np.frombuffer(data, dtype=np.float32).reshape(M, N)

	# Mark weight as cached on server
	if not a_only:
	_warm_keys[weight_key] = (K, N)

	return C

	def vk_matmul_plain(M, N, K, A, B):
	"""Plain (uncached) matmul for comparison."""
	hdr = struct.pack('<IIIIIIII', MAGIC, M, N, K, TYPE_PLAIN, 0, 0, 0)
	_send_all(hdr)
	_send_all(A.astype(np.float32).tobytes())
	_send_all(B.astype(np.float32).tobytes())
	rhdr = _recv_all(16)
	magic, status, rM, rN = struct.unpack('<IIII', rhdr)
	assert magic == MAGIC and status == 0
	return np.frombuffer(_recv_all(MN4), dtype=np.float32).reshape(M, N)

	# ─── Tests ────────────────────────────────────────────────────

	print(f"=== Vulkan Matmul Server v2 Test — VRAM Cache ===")
	print(f"Connecting to {HOST}:{PORT}...\n")

	print("── 1. Correctness (plain vs cached) ──")
	for M, N, K, label in [
	(1, 256, 256, "tiny"),
	(1, 1024, 1024, "1K gen"),
	(128, 2048, 2048, "128 prefill"),
	(1, 2048, 7168, "MoE expert"),
	]:
	A = np.random.randn(M, K).astype(np.float32)
	B = np.random.randn(K, N).astype(np.float32)
	key = _next_key()

	C_np = A @ B
	C_vk = vk_matmul_cached(M, N, K, A, B, key)
	err = np.max(np.abs(C_vk - C_np))
	print(f" {'✓' if err < 0.01 else '✗'} [{M}×{N}×{K}] {label:20s} maxerr={err:.2e}")

	print()
	print("── 2. Cache speedup — simulating 120B generation tokens ──")
	print(" (Upload weight ONCE, then send only activations for each token)")
	print()

	shapes = [
	(7168, 7168, "q_proj/k_proj/v_proj (attention)"),
	(7168, 7168, "o_proj (attention)"),
	(7168, 2048, "expert gate_proj (MoE FFN)"),
	(7168, 2048, "expert up_proj (MoE FFN)"),
	(2048, 7168, "expert down_proj (MoE FFN)"),
	]
	N_TOKENS = 20

	for (K, N, label) in shapes:
	B = np.random.randn(K, N).astype(np.float32)
	key = _next_key() # stable unique key — no id() address-reuse collisions
	# Remove from warm cache so first call uploads (key is always fresh here)
	_warm_keys.pop(key, None)

	cold_times = []
	warm_times = []

	for tok in range(N_TOKENS):
	A = np.random.randn(1, K).astype(np.float32) # single token
	t0 = time.time()
	C = vk_matmul_cached(1, N, K, A, B, key)
	elapsed = (time.time() - t0) * 1000
	if tok == 0:
	cold_times.append(elapsed) # first call = upload + compute
	else:
	warm_times.append(elapsed) # subsequent = activation only + compute

	cold_avg = sum(cold_times)
	warm_avg = sum(warm_times) / len(warm_times)
	speedup = cold_avg / warm_avg if warm_avg > 0 else 0
	wt_mb = K * N * 4 / 1e6

	print(f" [{K}×{N}] {label}")
	print(f" Cold (upload {wt_mb:.0f}MB weight): {cold_avg:.0f}ms")
	print(f" Warm (send 4KB activation only): {warm_avg:.1f}ms ← {speedup:.0f}× faster")
	print()

	print()
	print("── 3. Throughput: tokens/sec with cached weights ──")
	# Simulate a full 120B MoE forward pass for one token
	# Real 120B MoE: ~80 layers, each with 4 attn projections + 2 active experts × 3 projections
	LAYERS = 4 # subset for timing
	ATTN_PROJS = 4 # Q,K,V,O
	EXPERTS = 2 # active experts per token
	EXP_PROJS = 3 # gate, up, down

	# Pre-create weight matrices (simulating model loaded in POWER8 RAM)
	attn_weights = [np.random.randn(7168, 7168).astype(np.float32) for _ in range(ATTN_PROJS)]
	exp_weights = [(np.random.randn(7168, 2048).astype(np.float32),
	np.random.randn(7168, 2048).astype(np.float32),
	np.random.randn(2048, 7168).astype(np.float32))
	for _ in range(EXPERTS)]

	# Warm up the cache (first forward pass)
	print(" Warming up weight cache (first token upload)...")
	t_warmup = time.time()
	A_dummy = np.random.randn(1, 7168).astype(np.float32)
	for layer in range(LAYERS):
	for wi, W in enumerate(attn_weights):
	key = (layer * 1000 + wi) & 0xFFFFFFFFFFFFFFFF
	_warm_keys.pop(key, None)
	vk_matmul_cached(1, 7168, 7168, A_dummy, W, key)
	for ei, (Wg, Wu, Wd) in enumerate(exp_weights):
	kg = (layer * 1000 + 100 + ei * 3) & 0xFFFFFFFFFFFFFFFF
	ku = (layer * 1000 + 101 + ei * 3) & 0xFFFFFFFFFFFFFFFF
	kd = (layer * 1000 + 102 + ei * 3) & 0xFFFFFFFFFFFFFFFF
	for k, W, K, N in [(kg,Wg,7168,2048),(ku,Wu,7168,2048),(kd,Wd,2048,7168)]:
	_warm_keys.pop(k, None)
	vk_matmul_cached(1, N, K, np.random.randn(1,K).astype(np.float32), W, k)
	warmup_ms = (time.time() - t_warmup) * 1000
	print(f" Warmup done in {warmup_ms:.0f}ms\n")

	# Now time N tokens with warm cache
	N_TIME_TOKENS = 10
	t_start = time.time()
	for tok in range(N_TIME_TOKENS):
	A = np.random.randn(1, 7168).astype(np.float32)
	for layer in range(LAYERS):
	for wi, W in enumerate(attn_weights):
	key = (layer * 1000 + wi) & 0xFFFFFFFFFFFFFFFF
	vk_matmul_cached(1, 7168, 7168, A, W, key)
	for ei, (Wg, Wu, Wd) in enumerate(exp_weights):
	kg = (layer * 1000 + 100 + ei * 3) & 0xFFFFFFFFFFFFFFFF
	ku = (layer * 1000 + 101 + ei * 3) & 0xFFFFFFFFFFFFFFFF
	kd = (layer * 1000 + 102 + ei * 3) & 0xFFFFFFFFFFFFFFFF
	vk_matmul_cached(1, 2048, 7168, np.random.randn(1,7168).astype(np.float32), Wg, kg)
	vk_matmul_cached(1, 2048, 7168, np.random.randn(1,7168).astype(np.float32), Wu, ku)
	vk_matmul_cached(1, 7168, 2048, np.random.randn(1,2048).astype(np.float32), Wd, kd)

	elapsed_s = time.time() - t_start
	ms_per_tok = elapsed_s * 1000 / N_TIME_TOKENS
	ops_per_tok = LAYERS * (ATTN_PROJS + EXPERTS * EXP_PROJS)

	print(f" {N_TIME_TOKENS} tokens × {LAYERS} layers × "
	f"({ATTN_PROJS} attn + {EXPERTS}×{EXP_PROJS} MoE) ops")
	print(f" Time: {elapsed_s*1000:.0f}ms total = {ms_per_tok:.1f}ms/token")
	print(f" Throughput: {1000/ms_per_tok:.2f} tokens/sec "
	f"({ops_per_tok} matmuls/token, {ms_per_tok/ops_per_tok:.1f}ms each)")
	print(f"\n Cache stats: {len(_warm_keys)} weights resident in VRAM")
	print(f" (Extrapolated full 80-layer 120B: "
	f"~{ms_per_tok * 80/LAYERS:.0f}ms/token = "
	f"~{1000/(ms_per_tok * 80/LAYERS):.2f} t/s)")

	_sock.close()
	print("\nAll done.")