# SFT (Supervised Fine-Tuning) submodule for kt-kernel
# SPDX-License-Identifier: Apache-2.0

"""
SFT training support for KT-Kernel MoE.

This submodule adds training capabilities (forward/backward, LoRA, autograd,
distributed) on top of the inference-only kt_kernel base package.

Additional dependencies beyond base kt_kernel: torch.nn, torch.distributed, peft (optional).
"""

from .config import KTActivationPolicy, KTConfig
from .backend import (
    FP8_BACKEND,
    FP8_SFT_METHOD,
    INT8_BACKEND,
    INT8_SFT_METHOD,
    RAWINT4_BACKEND,
    RAWINT4_GROUP_SIZE,
    RAWINT4_SFT_METHOD,
    RAWINT4_WEIGHT_LAYOUT,
    get_fp8_runtime,
    get_int8_runtime,
    get_rawint4_checkpoint_contract,
    get_rawint4_runtime,
)
from .weight_manifest import validate_persistent_int8_weights
from .artifacts import (
    FUSED_EXPERT_LORA_NAME,
    KT_ADAPTER_MANIFEST_NAME,
    KT_NON_EXPERT_MANIFEST_NAME,
    KTAdapterManifest,
    KTArtifactError,
    KTNonExpertCache,
    KTPretrainedLoadPlan,
    claim_kt_routed_expert_subtrees,
    hide_kt_int8_routed_experts_from_dispatch,
    hide_kt_routed_experts_from_dispatch,
    is_kt_int8_routed_expert_base_parameter,
    is_kt_routed_expert_base_parameter,
    is_kt_routed_expert_parameter_name,
    is_kt_supported_moe_model,
    get_kt_fused_lora_exclude_modules,
    load_kt_adapter_artifacts,
    mark_kt_int8_routed_expert_base_parameters,
    prepare_kt_int8_non_expert_device_map,
    prepare_kt_non_expert_device_map,
    project_kt_int8_routed_experts_out_of_device_map,
    project_kt_routed_experts_out_of_device_map,
    resolve_kt_pretrained_artifacts,
    save_kt_adapter_artifacts,
    should_disable_kt_source_quantizer,
    validate_kt_prequantized_loading_info,
    validate_kt_pretrained_load,
    write_kt_non_expert_cache_manifest,
)
from .dist_utils import get_activation_checkpoint_context_fn
from .base import BaseSFTMoEWrapper, KExpertsSFTBuffer
from .amx import AMXSFTMoEWrapper
from .arch import (
    MOEArchConfig,
    get_moe_arch_config,
    get_moe_module,
    move_non_experts_to_gpu,
    get_expert_device,
    KTAMXError,
    KTAMXNotAvailableError,
    KTAMXModelNotSupportedError,
    KTAMXConfigError,
)
from .autograd import KTMoEFunction
from .conv3d_compat import is_vlm_conv3d_compatible, patch_vlm_conv3d
from .layer import KTMoELayerWrapper
from .weights import (
    extract_moe_weights,
    get_kt_expert_placeholders,
    load_block_fp8_experts_from_checkpoint_files,
    load_experts_from_checkpoint_files,
    load_experts_from_kt_weight_path,
    load_rawint4_experts_from_checkpoint_files,
    INT8ExpertWeights,
    BlockFP8ExpertWeights,
    RAWINT4ExpertWeights,
)
from .lora import (
    KTAdaptationResult,
    adapt_peft_lora,
    kt_adapt_peft_lora,
    get_kt_lora_named_params,
    get_kt_lora_params,
    get_kt_named_trainable_params,
    get_kt_rank_local_parameter_names,
    get_kt_trainable_params,
    update_kt_lora_pointers,
    sync_kt_lora_gradients,
    save_lora_experts_to_adapter,
    save_kt_moe_to_adapter,
    load_lora_experts_from_adapter,
    load_kt_moe_from_adapter,
    LoRAExpertMLP,
    LoRAExperts,
)
from .wrapper import (
    wrap_moe_layers_with_kt_wrapper,
    build_kt_device_map,
    build_kt_device_map_simplified,
    get_kt_loading_kwargs,
    load_kt_model,
)
from .profiler import collect_kt_sft_profile, format_kt_sft_profile, reset_kt_sft_profile

__all__ = [
    "KTConfig",
    "KTActivationPolicy",
    "INT8_BACKEND",
    "INT8_SFT_METHOD",
    "get_int8_runtime",
    "FP8_BACKEND",
    "FP8_SFT_METHOD",
    "get_fp8_runtime",
    "RAWINT4_BACKEND",
    "RAWINT4_GROUP_SIZE",
    "RAWINT4_SFT_METHOD",
    "RAWINT4_WEIGHT_LAYOUT",
    "get_rawint4_checkpoint_contract",
    "get_rawint4_runtime",
    "validate_persistent_int8_weights",
    "FUSED_EXPERT_LORA_NAME",
    "KT_ADAPTER_MANIFEST_NAME",
    "KT_NON_EXPERT_MANIFEST_NAME",
    "KTAdapterManifest",
    "KTArtifactError",
    "KTNonExpertCache",
    "KTPretrainedLoadPlan",
    "claim_kt_routed_expert_subtrees",
    "hide_kt_int8_routed_experts_from_dispatch",
    "hide_kt_routed_experts_from_dispatch",
    "is_kt_int8_routed_expert_base_parameter",
    "is_kt_routed_expert_base_parameter",
    "is_kt_routed_expert_parameter_name",
    "is_kt_supported_moe_model",
    "get_kt_fused_lora_exclude_modules",
    "load_kt_adapter_artifacts",
    "mark_kt_int8_routed_expert_base_parameters",
    "prepare_kt_int8_non_expert_device_map",
    "prepare_kt_non_expert_device_map",
    "project_kt_int8_routed_experts_out_of_device_map",
    "project_kt_routed_experts_out_of_device_map",
    "resolve_kt_pretrained_artifacts",
    "save_kt_adapter_artifacts",
    "should_disable_kt_source_quantizer",
    "validate_kt_prequantized_loading_info",
    "validate_kt_pretrained_load",
    "write_kt_non_expert_cache_manifest",
    "get_activation_checkpoint_context_fn",
    "BaseSFTMoEWrapper",
    "KExpertsSFTBuffer",
    "AMXSFTMoEWrapper",
    "MOEArchConfig",
    "get_moe_arch_config",
    "get_moe_module",
    "move_non_experts_to_gpu",
    "get_expert_device",
    "KTAMXError",
    "KTAMXNotAvailableError",
    "KTAMXModelNotSupportedError",
    "KTAMXConfigError",
    "KTMoEFunction",
    "patch_vlm_conv3d",
    "is_vlm_conv3d_compatible",
    "KTMoELayerWrapper",
    "extract_moe_weights",
    "get_kt_expert_placeholders",
    "load_block_fp8_experts_from_checkpoint_files",
    "load_experts_from_checkpoint_files",
    "load_experts_from_kt_weight_path",
    "load_rawint4_experts_from_checkpoint_files",
    "INT8ExpertWeights",
    "BlockFP8ExpertWeights",
    "RAWINT4ExpertWeights",
    "KTAdaptationResult",
    "adapt_peft_lora",
    "kt_adapt_peft_lora",
    "get_kt_lora_named_params",
    "get_kt_lora_params",
    "get_kt_named_trainable_params",
    "get_kt_rank_local_parameter_names",
    "get_kt_trainable_params",
    "update_kt_lora_pointers",
    "sync_kt_lora_gradients",
    "save_lora_experts_to_adapter",
    "save_kt_moe_to_adapter",
    "load_lora_experts_from_adapter",
    "load_kt_moe_from_adapter",
    "LoRAExpertMLP",
    "LoRAExperts",
    "wrap_moe_layers_with_kt_wrapper",
    "build_kt_device_map",
    "build_kt_device_map_simplified",
    "get_kt_loading_kwargs",
    "load_kt_model",
    "collect_kt_sft_profile",
    "format_kt_sft_profile",
    "reset_kt_sft_profile",
]