mirror of
https://github.com/kvcache-ai/ktransformers.git
synced 2026-04-29 12:19:50 +00:00
parent
4b235cdaa4
commit
9539ab91eb
4 changed files with 382 additions and 151 deletions
|
|
@ -4,6 +4,8 @@ Doctor command for kt-cli.
|
|||
Diagnoses environment issues and provides recommendations.
|
||||
"""
|
||||
|
||||
import glob
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
|
@ -29,6 +31,67 @@ from kt_kernel.cli.utils.environment import (
|
|||
)
|
||||
|
||||
|
||||
def _get_kt_kernel_info() -> dict:
|
||||
"""Get kt-kernel installation information."""
|
||||
info = {
|
||||
"installed": False,
|
||||
"version": None,
|
||||
"cpu_variant": None,
|
||||
"install_path": None,
|
||||
"available_variants": [],
|
||||
"extension_file": None,
|
||||
}
|
||||
|
||||
try:
|
||||
import kt_kernel
|
||||
|
||||
info["installed"] = True
|
||||
info["version"] = getattr(kt_kernel, "__version__", "unknown")
|
||||
info["cpu_variant"] = getattr(kt_kernel, "__cpu_variant__", "unknown")
|
||||
|
||||
# Get installation path
|
||||
info["install_path"] = os.path.dirname(kt_kernel.__file__)
|
||||
|
||||
# Find available .so files
|
||||
kt_kernel_dir = info["install_path"]
|
||||
so_files = glob.glob(os.path.join(kt_kernel_dir, "_kt_kernel_ext_*.so"))
|
||||
so_files.extend(glob.glob(os.path.join(kt_kernel_dir, "kt_kernel_ext*.so")))
|
||||
|
||||
# Parse variant names from filenames
|
||||
variants = set()
|
||||
for so_file in so_files:
|
||||
basename = os.path.basename(so_file)
|
||||
if "_kt_kernel_ext_" in basename:
|
||||
# Extract variant from _kt_kernel_ext_amx.cpython-311-x86_64-linux-gnu.so
|
||||
parts = basename.split("_")
|
||||
if len(parts) >= 4:
|
||||
variant = parts[3] # "amx" from "_kt_kernel_ext_amx..."
|
||||
if variant.startswith("avx"):
|
||||
# Normalize avx variants
|
||||
if variant in ["avx512", "avx512_bf16", "avx512_vbmi", "avx512_vnni", "avx512_base"]:
|
||||
variants.add("avx512")
|
||||
else:
|
||||
variants.add(variant)
|
||||
else:
|
||||
variants.add(variant)
|
||||
elif "kt_kernel_ext" in basename:
|
||||
variants.add("default")
|
||||
|
||||
info["available_variants"] = sorted(list(variants))
|
||||
|
||||
# Get current extension file
|
||||
if hasattr(kt_kernel, "kt_kernel_ext"):
|
||||
ext_module = kt_kernel.kt_kernel_ext
|
||||
info["extension_file"] = getattr(ext_module, "__file__", None)
|
||||
|
||||
except ImportError:
|
||||
info["installed"] = False
|
||||
except Exception as e:
|
||||
info["error"] = str(e)
|
||||
|
||||
return info
|
||||
|
||||
|
||||
def doctor(
|
||||
verbose: bool = typer.Option(False, "--verbose", "-v", help="Show detailed diagnostics"),
|
||||
) -> None:
|
||||
|
|
@ -157,6 +220,76 @@ def doctor(
|
|||
}
|
||||
)
|
||||
|
||||
# 6b. kt-kernel installation check
|
||||
kt_info = _get_kt_kernel_info()
|
||||
|
||||
if kt_info["installed"]:
|
||||
# Build display string for kt-kernel
|
||||
variant = kt_info["cpu_variant"]
|
||||
version = kt_info["version"]
|
||||
available_variants = kt_info["available_variants"]
|
||||
|
||||
# Determine status based on CPU variant
|
||||
if variant == "amx":
|
||||
kt_status = "ok"
|
||||
kt_hint = "AMX variant loaded - optimal performance"
|
||||
elif variant.startswith("avx512"):
|
||||
kt_status = "ok"
|
||||
kt_hint = "AVX512 variant loaded - good performance"
|
||||
elif variant == "avx2":
|
||||
kt_status = "warning"
|
||||
kt_hint = "AVX2 variant - consider upgrading CPU for AMX/AVX512"
|
||||
else:
|
||||
kt_status = "warning"
|
||||
kt_hint = f"Unknown variant: {variant}"
|
||||
|
||||
kt_value = f"v{version} ({variant.upper()})"
|
||||
if verbose and available_variants:
|
||||
kt_value += f" [dim] - available: {', '.join(available_variants)}[/dim]"
|
||||
|
||||
checks.append(
|
||||
{
|
||||
"name": "kt-kernel",
|
||||
"status": kt_status,
|
||||
"value": kt_value,
|
||||
"hint": kt_hint,
|
||||
}
|
||||
)
|
||||
|
||||
# Show extension file path in verbose mode
|
||||
if verbose and kt_info.get("extension_file"):
|
||||
ext_file = os.path.basename(kt_info["extension_file"])
|
||||
checks.append(
|
||||
{
|
||||
"name": " └─ Extension",
|
||||
"status": "ok",
|
||||
"value": ext_file,
|
||||
"hint": None,
|
||||
}
|
||||
)
|
||||
|
||||
# Show installation path in verbose mode
|
||||
if verbose and kt_info.get("install_path"):
|
||||
checks.append(
|
||||
{
|
||||
"name": " └─ Path",
|
||||
"status": "ok",
|
||||
"value": kt_info["install_path"],
|
||||
"hint": None,
|
||||
}
|
||||
)
|
||||
else:
|
||||
error_msg = kt_info.get("error", "Not installed")
|
||||
checks.append(
|
||||
{
|
||||
"name": "kt-kernel",
|
||||
"status": "error",
|
||||
"value": error_msg,
|
||||
"hint": "kt-kernel is required - run: pip install kt-kernel",
|
||||
}
|
||||
)
|
||||
issues_found = True
|
||||
|
||||
# 7. System memory (with frequency if available)
|
||||
mem_info = detect_memory_info()
|
||||
if mem_info.frequency_mhz and mem_info.type:
|
||||
|
|
@ -204,7 +337,6 @@ def doctor(
|
|||
# 6. Required packages
|
||||
packages = [
|
||||
("kt-kernel", ">=0.4.0", False), # name, version_req, required
|
||||
("ktransformers", ">=0.4.0", False),
|
||||
("sglang", ">=0.4.0", False),
|
||||
("torch", ">=2.4.0", True),
|
||||
("transformers", ">=4.45.0", True),
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue