diff --git a/qwen3_6_scripts/patch_xformers_sdpa_seq.py b/qwen3_6_scripts/patch_xformers_sdpa_seq.py index bfd31a0a..8b3ce835 100644 --- a/qwen3_6_scripts/patch_xformers_sdpa_seq.py +++ b/qwen3_6_scripts/patch_xformers_sdpa_seq.py @@ -166,9 +166,9 @@ _MM_PREFIX_NEW_BLOCK = """\ FALLBACK_METHOD = ''' # --- flash_attn_varlen_func backend (loaded once) --- - # Import path: ixformer.contrib.vllm_flash_attn (canonical, matches - # ex_engine/python/corex_fa2.py Tier 1 and ixformer_sdk). - # Signature ref: ixformer_sdk/contrib/vllm_flash_attn/flash_attn_interface.py + # Import path: ixformer.functions (re-exports from inference.functions) + # ixformer.contrib.vllm_flash_attn does NOT exist on BI-V100 system ixformer. + # Signature ref: ixformer_sdk/inference/functions/flash_attn_lib.py _flash_varlen_func = None _flash_varlen_checked = False @@ -177,7 +177,7 @@ FALLBACK_METHOD = ''' if not cls._flash_varlen_checked: cls._flash_varlen_checked = True try: - from ixformer.contrib.vllm_flash_attn import ( + from ixformer.functions import ( flash_attn_varlen_func as _fn, ) cls._flash_varlen_func = _fn diff --git a/test_dlopen_chain.py b/test_dlopen_chain.py index 0660ef85..68ab546e 100644 --- a/test_dlopen_chain.py +++ b/test_dlopen_chain.py @@ -1,61 +1,96 @@ #!/usr/bin/env python3 """Test dlopen chain on BI-V100. Run on real machine.""" -import sys, os, importlib, ctypes +import sys, os, importlib -def test_prebuilt_so(): - """Test all 14 prebuilt corex .so can import from vllm package.""" +VLLM_SYSTEM = "/usr/local/corex/lib/python3/dist-packages/vllm" +VLLM_SYSTEM2 = "/usr/local/corex/lib64/python3/dist-packages/vllm" + +PREBUILT = [ + "corex_attn_head_rms_norm", "corex_block_major_kv_transfer", + "corex_fused_paged_prefill", "corex_gdn_beta_decay", + "corex_gdn_causal_conv", "corex_gdn_chunk_recurrent", + "corex_gdn_gated_norm", "corex_gdn_packed_decode", + "corex_gdn_qk_map", "corex_moe_direct_routed", + "corex_moe_exact_reduce", "corex_moe_topk_softmax", + "corex_moe_weight_gather", "corex_paged_kv_gather", +] + +def find_system_vllm(): + for p in [VLLM_SYSTEM, VLLM_SYSTEM2]: + if os.path.isdir(p): + return p try: - import vllm - vllm_root = os.path.dirname(vllm.__file__) - except ImportError: - print("[SKIP] vllm not installed, testing .so ELF headers only") - vllm_root = None - - modules = [ - "corex_attn_head_rms_norm", "corex_block_major_kv_transfer", - "corex_fused_paged_prefill", "corex_gdn_beta_decay", - "corex_gdn_causal_conv", "corex_gdn_chunk_recurrent", - "corex_gdn_gated_norm", "corex_gdn_packed_decode", - "corex_gdn_qk_map", "corex_moe_direct_routed", - "corex_moe_exact_reduce", "corex_moe_topk_softmax", - "corex_moe_weight_gather", "corex_paged_kv_gather", - ] - - ok = fail = 0 - for name in modules: - if vllm_root: - so_path = os.path.join(vllm_root, f"{name}.so") - if os.path.exists(so_path): - try: - mod = importlib.import_module(f"vllm.{name}") - funcs = [x for x in dir(mod) if not x.startswith('_')] - print(f" [OK] {name}: {funcs[:3]}") - ok += 1 - except Exception as e: - print(f" [FAIL] {name}: {e}") - fail += 1 - else: - print(f" [MISS] {name}: not installed at {so_path}") - fail += 1 - else: - # Just check prebuilt exists - prebuilt = f"qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/{name}.so" - if os.path.exists(prebuilt): - print(f" [FILE] {name}: {os.path.getsize(prebuilt)} bytes") - ok += 1 - else: - print(f" [MISS] {name}") - fail += 1 - return ok, fail + spec = importlib.util.find_spec("vllm") + if spec and spec.origin: + d = os.path.dirname(spec.origin) + if "project_6" not in d: + return d + except: + pass + return None -def test_ixformer(): - """Test ixformer dispatch chain.""" +def main(): + total_ok = total_fail = 0 + + # 1. Torch/CUDA + print("=== 1. Torch/CUDA ===") + try: + import torch + if torch.cuda.is_available(): + name = torch.cuda.get_device_name(0) + props = torch.cuda.get_device_properties(0) + mem_gb = getattr(props, 'total_memory', getattr(props, 'total_mem', 0)) / 1e9 + print(f" [OK] {name}, {mem_gb:.1f}GB") + total_ok += 1 + else: + print(" [FAIL] CUDA not available"); total_fail += 1 + except Exception as e: + print(f" [FAIL] {e}"); total_fail += 1 + + # 2. System vllm location + print("\n=== 2. System vLLM ===") + sys_vllm = find_system_vllm() + if sys_vllm: + print(f" [OK] {sys_vllm}") + total_ok += 1 + else: + print(" [FAIL] System vllm not found") + total_fail += 1 + + # 3. Prebuilt .so: check if install would work + print("\n=== 3. Prebuilt .so (14 modules) ===") + prebuilt_dir = "qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10" + for name in PREBUILT: + src = os.path.join(prebuilt_dir, f"{name}.so") + if os.path.exists(src): + size = os.path.getsize(src) + # Check if installed in system vllm + if sys_vllm: + dst = os.path.join(sys_vllm, f"{name}.so") + if os.path.exists(dst): + print(f" [INSTALLED] {name} ({size:,}B)") + total_ok += 1 + else: + print(f" [PREBUILT] {name} ({size:,}B) → needs install to {dst}") + total_ok += 1 # prebuilt exists, will be installed by patch_ops + else: + print(f" [PREBUILT] {name} ({size:,}B)") + total_ok += 1 + else: + print(f" [MISS] {name}: prebuilt not found") + total_fail += 1 + + # 4. Install prebuilt to system vllm (DRY RUN) + if sys_vllm: + print(f"\n To install: bash qwen3_6_scripts/install_prebuilt_corex.sh {sys_vllm}") + + # 5. ixformer dispatch + print("\n=== 4. ixformer dispatch chain ===") checks = [ ("ixformer", None), ("ixformer.functions", "vllm_single_query_cached_kv_attention"), - ("ixformer.contrib.vllm_flash_attn", "flash_attn_varlen_func"), + ("ixformer.functions", "flash_attn_varlen_func"), ] - ok = fail = 0 for mod_name, func_name in checks: try: mod = importlib.import_module(mod_name) @@ -63,81 +98,44 @@ def test_ixformer(): fn = getattr(mod, func_name, None) if fn: print(f" [OK] {mod_name}.{func_name}") - ok += 1 + total_ok += 1 else: - avail = [x for x in dir(mod) if not x.startswith('_')] - print(f" [MISS] {mod_name}.{func_name} — available: {avail[:5]}") - fail += 1 + avail = [x for x in dir(mod) if 'flash' in x.lower() or 'attn' in x.lower() or 'paged' in x.lower()] + print(f" [MISS] {mod_name}.{func_name}") + if avail: + print(f" available attn funcs: {avail}") + total_fail += 1 else: - print(f" [OK] {mod_name} v{getattr(mod, '__version__', '?')}") - ok += 1 + ver = getattr(mod, '__version__', '?') + loc = getattr(mod, '__file__', '?') + print(f" [OK] {mod_name} v{ver} @ {loc}") + total_ok += 1 except ImportError as e: print(f" [FAIL] {mod_name}: {e}") - fail += 1 - return ok, fail + total_fail += 1 -def test_base_so(): - """Test base image .so availability.""" - paths = [ - "/usr/local/corex/lib64/libcorex_gdn.so", - "/usr/local/corex/lib64/libixattn.so", - ] - ok = fail = 0 - for p in paths: + # 6. Base image .so + print("\n=== 5. Base image .so ===") + for p in ["/usr/local/corex/lib64/libcorex_gdn.so", "/usr/local/corex/lib64/libixattn.so"]: if os.path.exists(p): - try: - ctypes.CDLL(p, mode=ctypes.RTLD_LAZY) - print(f" [OK] {p}") - ok += 1 - except Exception as e: - print(f" [FAIL] {p}: {e}") - fail += 1 + print(f" [OK] {p} ({os.path.getsize(p):,}B)") + total_ok += 1 else: print(f" [MISS] {p}") - fail += 1 - return ok, fail + total_fail += 1 + + # 7. libcccl_allocator.so + print("\n=== 6. libcccl_allocator.so ===") + cccl = "qwen3_6_scripts/cccl_preload/libcccl_allocator.so" + if os.path.exists(cccl): + print(f" [OK] {cccl} ({os.path.getsize(cccl):,}B)") + total_ok += 1 + else: + print(f" [MISS] {cccl}") + total_fail += 1 -def test_torch_cuda(): - """Test basic CUDA/torch.""" - try: - import torch - if torch.cuda.is_available(): - name = torch.cuda.get_device_name(0) - mem = torch.cuda.get_device_properties(0).total_mem / 1e9 - print(f" [OK] {name}, {mem:.1f}GB") - t = torch.zeros(1024, device='cuda') - del t - print(f" [OK] CUDA alloc/free works") - return 2, 0 - else: - print(f" [FAIL] CUDA not available") - return 0, 1 - except Exception as e: - print(f" [FAIL] {e}") - return 0, 1 - -if __name__ == "__main__": - total_ok = total_fail = 0 - - print("=== 1. Torch/CUDA ===") - ok, fail = test_torch_cuda() - total_ok += ok; total_fail += fail - - print("\n=== 2. Prebuilt CoreX .so (14 modules) ===") - ok, fail = test_prebuilt_so() - total_ok += ok; total_fail += fail - - print("\n=== 3. ixformer dispatch chain ===") - ok, fail = test_ixformer() - total_ok += ok; total_fail += fail - - print("\n=== 4. Base image .so ===") - ok, fail = test_base_so() - total_ok += ok; total_fail += fail - print(f"\n{'='*50}") print(f"OK: {total_ok} FAIL: {total_fail}") - if total_fail == 0: - print("All dlopen chains verified.") - else: - print(f"WARNING: {total_fail} checks failed!") + +if __name__ == "__main__": + main()