diff --git a/audit_so_usage.sh b/audit_so_usage.sh new file mode 100644 index 00000000..71739c5e --- /dev/null +++ b/audit_so_usage.sh @@ -0,0 +1,106 @@ +#!/bin/bash +set -euo pipefail +cat << 'PYEOF' | CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" python3 -u - +"""Audit: which .so functions are actually called in the hot path vs available but unused.""" +import importlib.util, os, sys + +SO_DIR = "qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10" +QWEN = "qwen3_6_scripts/qwen3_5.py" +PATCH = "qwen3_6_scripts/ex_engine/python/patch_vllm_hot_path.py" +XLLM_OPS = "qwen3_6_scripts/ex_engine/python/xllm_ops.py" + +# 1. Collect all exported functions from all .so +print("=" * 70) +print(" AUDIT: .so function usage") +print("=" * 70) + +so_exports = {} +for f in sorted(os.listdir(SO_DIR)): + if not f.endswith(".so"): + continue + name = f[:-3] + path = os.path.join(SO_DIR, f) + try: + spec = importlib.util.spec_from_file_location(name, path) + m = importlib.util.module_from_spec(spec) + spec.loader.exec_module(m) + fns = [x for x in dir(m) if not x.startswith("_")] + so_exports[name] = fns + except Exception as e: + so_exports[name] = [f"LOAD_ERROR: {e}"] + +# 2. Search for usage in qwen3_5.py, patch_vllm_hot_path.py, xllm_ops.py +code_files = {} +for label, path in [("qwen3_5.py", QWEN), ("patch_hot_path.py", PATCH), ("xllm_ops.py", XLLM_OPS)]: + try: + with open(path) as f: + code_files[label] = f.read() + except: + code_files[label] = "" + +# Also scan all ex_engine python files +for f in os.listdir("qwen3_6_scripts/ex_engine/python"): + if f.endswith(".py"): + path = os.path.join("qwen3_6_scripts/ex_engine/python", f) + try: + with open(path) as fh: + code_files[f"ex_engine/{f}"] = fh.read() + except: + pass + +all_code = "\n".join(code_files.values()) + +# 3. For each .so and function, check if it's referenced +print(f"\n{'SO Module':<35} {'Function':<30} {'Used?':<6} {'Where'}") +print("-" * 110) + +total_fns = 0 +used_fns = 0 +unused = [] + +for so_name in sorted(so_exports.keys()): + fns = so_exports[so_name] + for fn in fns: + if "LOAD_ERROR" in fn: + print(f"{so_name:<35} {fn}") + continue + total_fns += 1 + + # Search patterns: module.fn, .fn(, "fn" + found_in = [] + for label, code in code_files.items(): + if f".{fn}" in code or f'"{fn}"' in code or f"'{fn}'" in code: + found_in.append(label) + + is_used = len(found_in) > 0 + if is_used: + used_fns += 1 + else: + unused.append((so_name, fn)) + + where = ", ".join(found_in[:3]) if found_in else "" + marker = " ✓" if is_used else " ✗" + print(f"{so_name:<35} {fn:<30} {marker:<6} {where}") + +print(f"\n{'=' * 70}") +print(f" TOTAL: {used_fns}/{total_fns} functions used") +print(f" UNUSED: {total_fns - used_fns} functions") +print(f"{'=' * 70}") + +if unused: + print(f"\n === UNUSED FUNCTIONS ===") + for so_name, fn in unused: + print(f" {so_name}.{fn}") + +# 4. Check which ixformer_torch_ext functions exist but aren't wrapped +print(f"\n === ixformer_torch_ext available but not in any bridge .so ===") +ix_fns = [ + "ixformer_linear", "ixformer_linear_ex", "ixformer_linear_allreduce", + "linear_i8w8o32", "quantized_linear_awq", "quantized_linear_gptq", + "quantized_linear_int8", "quantized_linear_float4", "ixformer_quantized_linear", + "silu_and_mul_forward", "rms_norm_forward", "fused_add_rms_norm_forward", +] +for fn in ix_fns: + in_bridge = fn in all_code + print(f" {fn:<40} {'✓ wrapped' if in_bridge else '✗ NOT wrapped'}") +PYEOF \ No newline at end of file diff --git a/bench_linear_patch.sh b/bench_linear_patch.sh new file mode 100644 index 00000000..b0da787b --- /dev/null +++ b/bench_linear_patch.sh @@ -0,0 +1,105 @@ +#!/bin/bash +set -euo pipefail +cat << 'PYEOF' | CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" python3 -u - +import torch, importlib.util, time +torch.cuda.set_device(0) +dev = torch.device("cuda:0") + +SO = "qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10" +def load_so(name): + spec = importlib.util.spec_from_file_location(name, f"{SO}/{name}.so") + m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m); return m + +bridge = load_so("ix_moe_bridge") +act_m = load_so("xllm_activation") + +H = 2048 + +def bench(name, fn, N=1000): + for _ in range(100): fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(N): fn() + torch.cuda.synchronize() + us = (time.perf_counter() - t0) / N * 1e6 + print(f" {name:50s}: {us:8.1f} us") + return us + +x = torch.randn(1, H, device=dev, dtype=torch.float16) + +# Match upstream gemv_conditions: m <= 1, k % 32 == 0, n % 2 == 0, no bias +# Test ALL linear ops in qwen3_5.py decode path + +print("=== Every linear op in one decode step (TP=4) ===") +print("--- Attention layer (32 layers) ---") + +# QKV: (1,2048) @ (1024,2048)^T → (1,1024) [heads*head_dim + 2*kv_heads*head_dim] +w_qkv = torch.randn(1024, H, device=dev, dtype=torch.float16) * 0.01 +bench("qkv F.linear (1,2048)→(1,1024)", lambda: torch.nn.functional.linear(x, w_qkv)) +bench("qkv bridge.linear", lambda: bridge.linear(x, w_qkv, None)) + +# O_proj: (1,768) @ (2048,768)^T → (1,2048) +x_o = torch.randn(1, 768, device=dev, dtype=torch.float16) +w_o = torch.randn(H, 768, device=dev, dtype=torch.float16) * 0.01 +bench("o_proj F.linear (1,768)→(1,2048)", lambda: torch.nn.functional.linear(x_o, w_o)) +bench("o_proj bridge.linear", lambda: bridge.linear(x_o, w_o, None)) + +print("\n--- GDN layer (4 layers) ---") +# GDN in_proj: (1,2048) @ (3852,2048)^T → (1,3852) +w_gdn = torch.randn(3852, H, device=dev, dtype=torch.float16) * 0.01 +bench("gdn_proj F.linear (1,2048)→(1,3852)", lambda: torch.nn.functional.linear(x, w_gdn)) +bench("gdn_proj bridge.linear", lambda: bridge.linear(x, w_gdn, None)) + +# GDN o_proj: (1,1536) @ (2048,1536)^T → (1,2048) +x_gdn_o = torch.randn(1, 1536, device=dev, dtype=torch.float16) +w_gdn_o = torch.randn(H, 1536, device=dev, dtype=torch.float16) * 0.01 +bench("gdn_oproj F.linear (1,1536)→(1,2048)", lambda: torch.nn.functional.linear(x_gdn_o, w_gdn_o)) +bench("gdn_oproj bridge.linear", lambda: bridge.linear(x_gdn_o, w_gdn_o, None)) + +print("\n--- MoE shared expert (36 layers) ---") +I_shared = 128 +w_gu = torch.randn(2*I_shared, H, device=dev, dtype=torch.float16) * 0.01 +w_down = torch.randn(H, I_shared, device=dev, dtype=torch.float16) * 0.01 +bench("shared gate_up F.linear (1,2048)→(1,256)", lambda: torch.nn.functional.linear(x, w_gu)) +bench("shared gate_up bridge.linear", lambda: bridge.linear(x, w_gu, None)) +x_down = torch.randn(1, I_shared, device=dev, dtype=torch.float16) +bench("shared down F.linear (1,128)→(1,2048)", lambda: torch.nn.functional.linear(x_down, w_down)) +bench("shared down bridge.linear", lambda: bridge.linear(x_down, w_down, None)) + +print("\n--- Router (36 layers) ---") +w_router = torch.randn(257, H, device=dev, dtype=torch.float16) * 0.01 +bench("router F.linear (1,2048)→(1,257)", lambda: torch.nn.functional.linear(x, w_router)) +bench("router bridge.linear", lambda: bridge.linear(x, w_router, None)) + +print("\n--- LM head (1x) ---") +w_lm = torch.randn(37984, H, device=dev, dtype=torch.float16) * 0.01 +bench("lm_head F.linear (1,2048)→(1,37984)", lambda: torch.nn.functional.linear(x, w_lm)) +bench("lm_head bridge.linear", lambda: bridge.linear(x, w_lm, None)) + +# === Total impact === +print("\n=== Projected total decode step savings ===") +shapes = [ + ("attn_qkv", 32, (1024, H)), + ("attn_o", 32, (H, 768)), + ("gdn_proj", 4, (3852, H)), + ("gdn_o", 4, (H, 1536)), + ("shared_gu", 36, (2*I_shared, H)), + ("shared_down",36, (H, I_shared)), + ("router", 36, (257, H)), + ("lm_head", 1, (37984, H)), +] +total_torch = 0 +total_bridge = 0 +for name, count, (N, K) in shapes: + w = torch.randn(N, K, device=dev, dtype=torch.float16) * 0.01 + xi = torch.randn(1, K, device=dev, dtype=torch.float16) + t_torch = bench(f" {name} F.linear", lambda xi=xi, w=w: torch.nn.functional.linear(xi, w), N=500) + t_bridge = bench(f" {name} bridge", lambda xi=xi, w=w: bridge.linear(xi, w, None), N=500) + total_torch += t_torch * count + total_bridge += t_bridge * count + speedup = t_torch / t_bridge if t_bridge > 0 else 0 + print(f" → x{count}: {t_torch*count:.0f} → {t_bridge*count:.0f} us ({speedup:.1f}x)") + +print(f"\n TOTAL linear ops: {total_torch:.0f} → {total_bridge:.0f} us") +print(f" Savings: {total_torch - total_bridge:.0f} us = {(total_torch-total_bridge)/1000:.1f} ms") +PYEOF \ No newline at end of file