chore: import upstream snapshot with attribution

kt-kernel/python/cli/utils/quant_interactive.py

@@ -0,0 +1,347 @@
+"""
+Interactive configuration for kt quant command.
+
+Provides rich, multi-step interactive configuration for model quantization.
+"""
+
+from typing import Optional, Dict, Any
+from pathlib import Path
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+from rich.prompt import Prompt, Confirm, IntPrompt
+from kt_kernel.cli.i18n import t
+
+
+console = Console()
+
+
+def select_model_to_quantize() -> Optional[Any]:
+    """Select model to quantize interactively."""
+    from kt_kernel.cli.utils.user_model_registry import UserModelRegistry
+    from kt_kernel.cli.commands.model import is_amx_weights, SHA256_STATUS_MAP
+    from kt_kernel.cli.utils.model_table_builder import build_moe_gpu_table
+
+    registry = UserModelRegistry()
+    all_models = registry.list_models()
+
+    # Filter MoE models only (safetensors, not AMX, is_moe=True)
+    quant_models = []
+    for model in all_models:
+        if model.format == "safetensors":
+            # Skip AMX models
+            is_amx, _ = is_amx_weights(model.path)
+            if is_amx:
+                continue
+
+            # Only include MoE models
+            if model.is_moe:
+                quant_models.append(model)
+
+    if not quant_models:
+        console.print(f"[yellow]{t('quant_no_moe_models')}[/yellow]")
+        console.print()
+        console.print(f"  {t('quant_only_moe')}")
+        console.print()
+        console.print(f"  {t('quant_add_models', command='kt model scan')}")
+        console.print(f"  {t('quant_add_models', command='kt model add <path>')}")
+        return None
+
+    # Display models
+    console.print()
+    console.print(f"[bold green]{t('quant_moe_available')}[/bold green]")
+    console.print()
+
+    # Use shared table builder
+    table, displayed_models = build_moe_gpu_table(
+        models=quant_models, status_map=SHA256_STATUS_MAP, show_index=True, start_index=1
+    )
+
+    console.print(table)
+    console.print()
+
+    choice = IntPrompt.ask(t("quant_select_model"), default=1, show_choices=False)
+
+    if choice < 1 or choice > len(displayed_models):
+        console.print(f"[red]{t('quant_invalid_choice')}[/red]")
+        return None
+
+    return displayed_models[choice - 1]
+
+
+def configure_quantization_method() -> Dict[str, str]:
+    """Select quantization method and input type."""
+    console.print()
+    console.print(Panel(f"[bold cyan]{t('quant_step2_method')}[/bold cyan]", expand=False))
+    console.print()
+
+    # Method selection
+    console.print(f"[bold]{t('quant_method_label')}[/bold]")
+    console.print(f"  [cyan][1][/cyan] {t('quant_int4_desc')}")
+    console.print(f"  [cyan][2][/cyan] {t('quant_int8_desc')}")
+    console.print()
+
+    method_choice = Prompt.ask(t("quant_select_method"), choices=["1", "2"], default="1")
+    method = "int4" if method_choice == "1" else "int8"
+
+    console.print()
+    console.print(f"[bold]{t('quant_input_type_label')}[/bold]")
+    console.print(f"  [cyan][1][/cyan] {t('quant_fp8_desc')}")
+    console.print(f"  [cyan][2][/cyan] {t('quant_fp16_desc')}")
+    console.print(f"  [cyan][3][/cyan] {t('quant_bf16_desc')}")
+    console.print()
+
+    input_choice = Prompt.ask(t("quant_select_input_type"), choices=["1", "2", "3"], default="1")
+    input_type_map = {"1": "fp8", "2": "fp16", "3": "bf16"}
+    input_type = input_type_map[input_choice]
+
+    return {"method": method, "input_type": input_type}
+
+
+def configure_cpu_params(max_cores: int, max_numa: int) -> Dict[str, Any]:
+    """Configure CPU parameters."""
+    console.print()
+    console.print(Panel(f"[bold cyan]{t('quant_step3_cpu')}[/bold cyan]", expand=False))
+    console.print()
+
+    def clamp(value: int, min_val: int, max_val: int, default: int) -> int:
+        """Clamp value to range or return default if out of bounds."""
+        if min_val <= value <= max_val:
+            return max(min_val, min(value, max_val))
+        return default
+
+    default_threads = int(max_cores * 0.8)
+    cpu_threads = IntPrompt.ask(t("quant_cpu_threads_prompt", max=max_cores), default=default_threads)
+    cpu_threads = clamp(cpu_threads, 1, max_cores, default_threads)
+
+    numa_nodes = IntPrompt.ask(t("quant_numa_nodes_prompt", max=max_numa), default=max_numa)
+    numa_nodes = clamp(numa_nodes, 1, max_numa, max_numa)
+
+    # Ask about GPU usage
+    console.print()
+    console.print(f"[bold]{t('quant_use_gpu_label')}[/bold]")
+    console.print(f"  [dim]{t('quant_gpu_speedup')}[/dim]")
+    console.print()
+    use_gpu = Confirm.ask(t("quant_enable_gpu"), default=True)
+
+    return {"cpu_threads": cpu_threads, "numa_nodes": numa_nodes, "use_gpu": use_gpu}
+
+
+def configure_output_path(model: Any, method: str, numa_nodes: int) -> Path:
+    """Configure output path for quantized weights."""
+    from kt_kernel.cli.config.settings import get_settings
+
+    console.print()
+    console.print(Panel(f"[bold cyan]{t('quant_step4_output')}[/bold cyan]", expand=False))
+    console.print()
+
+    # Generate default output path
+    model_path = Path(model.path)
+    method_upper = method.upper()
+    settings = get_settings()
+
+    # Priority: paths.weights > paths.models[0] > model's parent directory
+    weights_dir = settings.weights_dir
+    if weights_dir and weights_dir.exists():
+        # Use configured weights directory (highest priority)
+        default_output = weights_dir / f"{model_path.name}-AMX{method_upper}-NUMA{numa_nodes}"
+    else:
+        # Use first model storage path
+        model_paths = settings.get_model_paths()
+        if model_paths and model_paths[0].exists():
+            default_output = model_paths[0] / f"{model_path.name}-AMX{method_upper}-NUMA{numa_nodes}"
+        else:
+            # Fallback to model's parent directory
+            default_output = model_path.parent / f"{model_path.name}-AMX{method_upper}-NUMA{numa_nodes}"
+
+    console.print(f"[dim]{t('quant_default_path')}[/dim]", default_output)
+    console.print()
+
+    use_default = Confirm.ask(t("quant_use_default"), default=True)
+
+    if use_default:
+        return default_output
+
+    custom_path = Prompt.ask(t("quant_custom_path"), default=str(default_output))
+
+    return Path(custom_path)
+
+
+def calculate_quantized_size(source_path: Path, input_type: str, quant_method: str) -> tuple[float, float]:
+    """
+    Calculate source model size and estimated quantized size.
+
+    Args:
+        source_path: Path to source model
+        input_type: Input type (fp8, fp16, bf16)
+        quant_method: Quantization method (int4, int8)
+
+    Returns:
+        Tuple of (source_size_gb, estimated_quant_size_gb)
+    """
+    # Calculate source model size
+    try:
+        total_bytes = sum(f.stat().st_size for f in source_path.glob("*.safetensors") if f.is_file())
+        source_size_gb = total_bytes / (1024**3)
+    except Exception:
+        return 0.0, 0.0
+
+    # Bits mapping
+    input_bits = {"fp8": 8, "fp16": 16, "bf16": 16}
+    quant_bits = {"int4": 4, "int8": 8}
+
+    input_bit = input_bits.get(input_type, 16)
+    quant_bit = quant_bits.get(quant_method, 4)
+
+    # Estimate: source_size * (quant_bits / input_bits)
+    ratio = quant_bit / input_bit
+    estimated_size_gb = source_size_gb * ratio
+
+    return source_size_gb, estimated_size_gb
+
+
+def check_disk_space(output_path: Path, required_size_gb: float) -> tuple[float, bool]:
+    """
+    Check available disk space at output path.
+
+    Args:
+        output_path: Target output path
+        required_size_gb: Required space in GB
+
+    Returns:
+        Tuple of (available_gb, is_sufficient)
+        is_sufficient is True if available >= required * 1.2
+    """
+    import shutil
+
+    try:
+        # Get parent directory that exists
+        check_path = output_path.parent if not output_path.exists() else output_path
+        while not check_path.exists() and check_path != check_path.parent:
+            check_path = check_path.parent
+
+        stat = shutil.disk_usage(check_path)
+        available_gb = stat.free / (1024**3)
+
+        # Check if available space >= required * 1.2 (20% buffer)
+        is_sufficient = available_gb >= (required_size_gb * 1.2)
+
+        return available_gb, is_sufficient
+    except Exception:
+        return 0.0, False
+
+
+def interactive_quant_config() -> Optional[Dict[str, Any]]:
+    """
+    Interactive configuration for kt quant.
+
+    Returns configuration dict or None if cancelled.
+    """
+    from kt_kernel.cli.utils.environment import detect_cpu_info
+
+    # Get CPU info
+    cpu_info = detect_cpu_info()
+
+    # Step 1: Select model
+    model = select_model_to_quantize()
+    if not model:
+        return None
+
+    # Step 1.5: Pre-quantization verification (optional)
+    from kt_kernel.cli.utils.user_model_registry import UserModelRegistry
+    from kt_kernel.cli.utils.model_verifier import pre_operation_verification
+
+    user_registry = UserModelRegistry()
+    user_model_obj = user_registry.find_by_path(model.path)
+
+    if user_model_obj and user_model_obj.format == "safetensors":
+        pre_operation_verification(user_model_obj, user_registry, operation_name="quantizing")
+
+    # Step 2: Configure quantization method
+    quant_config = configure_quantization_method()
+
+    # Step 3: Configure CPU parameters
+    cpu_config = configure_cpu_params(cpu_info.threads, cpu_info.numa_nodes)  # Use logical threads
+
+    # Step 4: Configure output path
+    output_path = configure_output_path(model, quant_config["method"], cpu_config["numa_nodes"])
+
+    # Step 4.5: Check if output path already exists and generate unique name
+    if output_path.exists():
+        console.print()
+        console.print(t("quant_output_exists_warn", path=str(output_path)))
+        console.print()
+
+        # Generate unique name by adding suffix
+        original_name = output_path.name
+        parent_dir = output_path.parent
+        counter = 2
+
+        while output_path.exists():
+            new_name = f"{original_name}-{counter}"
+            output_path = parent_dir / new_name
+            counter += 1
+
+        console.print(t("quant_using_unique_name", path=str(output_path)))
+        console.print()
+
+    # Step 5: Calculate space requirements and check availability
+    console.print()
+    console.print(Panel(f"[bold cyan]{t('quant_disk_analysis')}[/bold cyan]", expand=False))
+    console.print()
+
+    source_size_gb, estimated_size_gb = calculate_quantized_size(
+        Path(model.path), quant_config["input_type"], quant_config["method"]
+    )
+
+    available_gb, is_sufficient = check_disk_space(output_path, estimated_size_gb)
+
+    console.print(f"  {t('quant_source_size'):<26} [cyan]{source_size_gb:.2f} GB[/cyan]")
+    console.print(f"  {t('quant_estimated_size'):<26} [yellow]{estimated_size_gb:.2f} GB[/yellow]")
+    console.print(
+        f"  {t('quant_available_space'):<26} [{'green' if is_sufficient else 'red'}]{available_gb:.2f} GB[/{'green' if is_sufficient else 'red'}]"
+    )
+    console.print()
+
+    if not is_sufficient:
+        required_with_buffer = estimated_size_gb * 1.2
+        console.print(f"[bold red]⚠ {t('quant_insufficient_space')}[/bold red]")
+        console.print()
+        console.print(f"  {t('quant_required_space'):<26} [yellow]{required_with_buffer:.2f} GB[/yellow]")
+        console.print(f"  {t('quant_available_space'):<26} [red]{available_gb:.2f} GB[/red]")
+        console.print(f"  {t('quant_shortage'):<26} [red]{required_with_buffer - available_gb:.2f} GB[/red]")
+        console.print()
+        console.print(f"  {t('quant_may_fail')}")
+        console.print()
+
+        if not Confirm.ask(f"[yellow]{t('quant_continue_anyway')}[/yellow]", default=False):
+            console.print(f"[yellow]{t('quant_cancelled')}[/yellow]")
+            return None
+        console.print()
+
+    # Summary and confirmation
+    console.print()
+    console.print(Panel(f"[bold cyan]{t('quant_config_summary')}[/bold cyan]", expand=False))
+    console.print()
+    console.print(f"  {t('quant_summary_model'):<15} {model.name}")
+    console.print(f"  {t('quant_summary_method'):<15} {quant_config['method'].upper()}")
+    console.print(f"  {t('quant_summary_input_type'):<15} {quant_config['input_type'].upper()}")
+    console.print(f"  {t('quant_summary_cpu_threads'):<15} {cpu_config['cpu_threads']}")
+    console.print(f"  {t('quant_summary_numa'):<15} {cpu_config['numa_nodes']}")
+    console.print(f"  {t('quant_summary_gpu'):<15} {t('yes') if cpu_config['use_gpu'] else t('no')}")
+    console.print(f"  {t('quant_summary_output'):<15} {output_path}")
+    console.print()
+
+    if not Confirm.ask(f"[bold green]{t('quant_start_question')}[/bold green]", default=True):
+        console.print(f"[yellow]{t('quant_cancelled')}[/yellow]")
+        return None
+
+    return {
+        "model": model,
+        "method": quant_config["method"],
+        "input_type": quant_config["input_type"],
+        "cpu_threads": cpu_config["cpu_threads"],
+        "numa_nodes": cpu_config["numa_nodes"],
+        "use_gpu": cpu_config["use_gpu"],
+        "output_path": output_path,
+    }