[refactor] Add KTransformers AMX MoE SFT support via Accelerate (#10430)

Co-authored-by: mrhaoxx <mr.haoxx@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml

@@ -0,0 +1,25 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: FSDP
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_offload_params: false
+  fsdp_reshard_after_forward: true
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_version: 2
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4 # Adjust based on your GPU count; 4 is suitable for 4 GPUs
+rdzv_backend: static
+same_network: true
+use_cpu: false
+
+kt_config:
+  enabled: true
+  kt_backend: AMXBF16  # Use with original BF16 expert weights.
+  kt_num_threads: 96
+  kt_tp_enabled: true
+  kt_threadpool_count: 2
+  kt_max_cache_depth: 2
+  kt_share_backward_bb: true
+  lora_rank: 8

examples/ktransformers/accelerate/fsdp2_kt_int4.yaml

@@ -0,0 +1,25 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: FSDP
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_offload_params: false
+  fsdp_reshard_after_forward: true
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_version: 2
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4 # Adjust based on your GPU count; 4 is suitable for 4 GPUs
+rdzv_backend: static
+same_network: true
+use_cpu: false
+
+kt_config:
+  enabled: true
+  kt_backend: AMXINT4  # Use with online-converted INT4 expert weights
+  kt_num_threads: 96
+  kt_tp_enabled: true
+  kt_threadpool_count: 2
+  kt_max_cache_depth: 2
+  kt_share_backward_bb: true
+  lora_rank: 8

examples/ktransformers/accelerate/fsdp2_kt_int8.yaml

@@ -0,0 +1,25 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: FSDP
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_offload_params: false
+  fsdp_reshard_after_forward: true
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_version: 2
+mixed_precision: bf16
+num_machines: 1
+num_processes: 4 # Adjust based on your GPU count; 4 is suitable for 4 GPUs
+rdzv_backend: static
+same_network: true
+use_cpu: false
+
+kt_config:
+  enabled: true
+  kt_backend: AMXINT8  # Use with online-converted INT8 expert weights
+  kt_num_threads: 96
+  kt_tp_enabled: true
+  kt_threadpool_count: 2
+  kt_max_cache_depth: 2
+  kt_share_backward_bb: true
+  lora_rank: 8

examples/ktransformers/accelerate/fsdp2_kt_int8_1gpu.yaml

@@ -0,0 +1,25 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: FSDP
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_offload_params: false
+  fsdp_reshard_after_forward: true
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_version: 2
+mixed_precision: bf16
+num_machines: 1
+num_processes: 1 # Adjust based on your GPU count; 1 is suitable for 1 GPU
+rdzv_backend: static
+same_network: true
+use_cpu: false
+
+kt_config:
+  enabled: true
+  kt_backend: AMXINT8  # Use with online-converted INT8 expert weights
+  kt_num_threads: 96
+  kt_tp_enabled: true
+  kt_threadpool_count: 2
+  kt_max_cache_depth: 2
+  kt_share_backward_bb: true
+  lora_rank: 8

examples/ktransformers/accelerate/fsdp2_kt_int8_8gpu.yaml

@@ -0,0 +1,25 @@
+compute_environment: LOCAL_MACHINE
+distributed_type: FSDP
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_offload_params: false
+  fsdp_reshard_after_forward: true
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_version: 2
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8 # Adjust based on your GPU count; 8 is suitable for 8 GPUs
+rdzv_backend: static
+same_network: true
+use_cpu: false
+
+kt_config:
+  enabled: true
+  kt_backend: AMXINT8  # Use with online-converted INT8 expert weights
+  kt_num_threads: 96
+  kt_tp_enabled: true
+  kt_threadpool_count: 2
+  kt_max_cache_depth: 2
+  kt_share_backward_bb: true
+  lora_rank: 8

examples/ktransformers/infer_lora/deepseek2_lora_sft_kt.yaml

@@ -1,10 +0,0 @@
-model_name_or_path: deepseek-ai/DeepSeek-V2-Lite
-adapter_name_or_path: saves/Kllama_deepseekV2
-template: deepseek
-infer_backend: ktransformers  # choices: [huggingface, vllm, sglang, ktransformers]
-trust_remote_code: true
-
-use_kt: true # use KTransformers as LoRA sft backend to inference
-kt_optimize_rule: examples/ktransformers/kt_optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx.yaml
-cpu_infer: 32
-chunk_size: 8192

examples/ktransformers/infer_lora/deepseek3_kt.yaml

@@ -1,9 +0,0 @@
-model_name_or_path: opensourcerelease/DeepSeek-V3-bf16
-template: deepseek3
-infer_backend: ktransformers  # choices: [huggingface, vllm, sglang, ktransformers]
-trust_remote_code: true
-
-use_kt: true # use KTransformers as LoRA sft backend to inference
-kt_optimize_rule: examples/ktransformers/kt_optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml
-cpu_infer: 32
-chunk_size: 8192

examples/ktransformers/infer_lora/deepseek3_lora_sft_kt.yaml

@@ -1,10 +0,0 @@
-model_name_or_path: opensourcerelease/DeepSeek-V3-bf16
-adapter_name_or_path: saves/Kllama_deepseekV3
-template: deepseek3
-infer_backend: ktransformers  # choices: [huggingface, vllm, sglang, ktransformers]
-trust_remote_code: true
-
-use_kt: true # use KTransformers as LoRA sft backend to inference
-kt_optimize_rule: examples/ktransformers/kt_optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml
-cpu_infer: 32
-chunk_size: 8192

examples/ktransformers/infer_lora/qwen3moe_lora_sft_kt.yaml

@@ -1,10 +0,0 @@
-model_name_or_path: Qwen/Qwen3-235B-A22B-Instruct-2507
-adapter_name_or_path: saves/Kllama_Qwen3MoE_235bA22b
-template: qwen3_nothink
-infer_backend: ktransformers  # choices: [huggingface, vllm, sglang, ktransformers]
-trust_remote_code: true
-
-use_kt: true # use KTransformers as LoRA sft backend to inference
-kt_optimize_rule: examples/ktransformers/kt_optimize_rules/Qwen3Moe-sft-amx.yaml
-cpu_infer: 32
-chunk_size: 8192

examples/ktransformers/kt_optimize_rules/DeepSeek-V2-Chat-sft-amx.yaml

@@ -1,69 +0,0 @@
-- match:
-    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^lm_head"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\..*\\.mlp$"
-    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model\\.layers\\..*\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cuda"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KSFTExpertsCPU"
-      out_device: "cuda"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False # don't recursively inject submodules of this module
-- match:
-    name: "^model\\.layers\\..*\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KDeepseekV2Model"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cpu"
-      prefill_device: "cpu"

examples/ktransformers/kt_optimize_rules/DeepSeek-V2-Chat.yaml

@@ -1,68 +0,0 @@
-- match:
-    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearMarlin"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^lm_head"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearMarlin"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\..*\\.mlp$"
-    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model\\.layers\\..*\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cuda"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KExpertsCPU"
-      out_device: "cuda"
-  recursive: False # don't recursively inject submodules of this module
-- match:
-    name: "^model\\.layers\\..*\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KDeepseekV2Model"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cpu"
-      prefill_device: "cpu"

examples/ktransformers/kt_optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx-multi-gpu.yaml

@@ -1,139 +0,0 @@
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-        generate_device: "cpu"
-        prefill_device: "cpu"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9])\\."
-    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-- match:
-    name: "^model\\.layers\\.([12][0-9])\\."
-    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9])\\.(?!.*self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\.([12][0-9])\\.(?!.*self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9])\\.mlp$"
-    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-- match:
-    name: "^model\\.layers\\.([12][0-9])\\.mlp$"
-    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9])\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cuda:0"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op:  "KSFTExpertsCPU"
-      out_device: "cuda:0"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False # don't recursively inject submodules of this module
-
-- match:
-    name: "^model\\.layers\\.([12][0-9])\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cuda:1"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op:  "KSFTExpertsCPU"
-      out_device: "cuda:1"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False # don't recursively inject submodules of this module
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9])\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-- match:
-    name: "^model\\.layers\\.([12][0-9])\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KDeepseekV2Model"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill
-      transfer_map:
-        10: "cuda:1"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9])\\."
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-
-- match:
-    name: "^lm_head"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "(^model\\.layers\\.([12][0-9])\\.)|(model.norm)"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"

examples/ktransformers/kt_optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx.yaml

@@ -1,69 +0,0 @@
-- match:
-    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^lm_head"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\..*\\.mlp$"
-    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model\\.layers\\..*\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cpu"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KSFTExpertsCPU"
-      out_device: "cuda"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False # don't recursively inject submodules of this module
-- match:
-    name: "^model\\.layers\\..*\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KDeepseekV2Model"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cpu"
-      prefill_device: "cpu"

examples/ktransformers/kt_optimize_rules/DeepSeek-V2-Lite-Chat-sft.yaml

@@ -1,68 +0,0 @@
-- match:
-    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^lm_head"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\..*\\.mlp$"
-    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model\\.layers\\..*\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cpu"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KSFTExpertsCPU"
-      out_device: "cuda"
-  recursive: False # don't recursively inject submodules of this module
-- match:
-    name: "^model\\.layers\\..*\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KDeepseekV2Model"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cpu"
-      prefill_device: "cpu"

examples/ktransformers/kt_optimize_rules/DeepSeek-V2-Lite-Chat.yaml

@@ -1,68 +0,0 @@
-- match:
-    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearMarlin"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^lm_head"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearMarlin"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\..*\\.mlp$"
-    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model\\.layers\\..*\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cuda"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KExpertsCPU"
-      out_device: "cuda"
-  recursive: False # don't recursively inject submodules of this module
-- match:
-    name: "^model\\.layers\\..*\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KDeepseekV2Model"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cpu"
-      prefill_device: "cpu"

examples/ktransformers/kt_optimize_rules/DeepSeek-V3-Chat-amx.yaml

@@ -1,77 +0,0 @@
-- match:
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-
-- match:
-    name: "^lm_head$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearMarlin"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearMarlin"
-      prefill_op: "KLinearTorch"
-- match:
-    name: "^model\\.layers\\..*\\.mlp$"
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV3MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    class: ktransformers.models.modeling_deepseek_v3.MoEGate
-  replace:
-    class: ktransformers.operators.gate.KMoEGate
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-- match:
-    name: "^model\\.layers\\..*\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cuda"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KExpertsCPU"
-      out_device: "cuda"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False # don't recursively inject submodules of this module
-- match:
-    name: "^model\\.layers\\..*\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      absorb_for_prefill: False # change this to True to enable long context(prefill may slower).
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KDeepseekV2Model"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cpu"
-      prefill_device: "cpu"

examples/ktransformers/kt_optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu-4.yaml

@@ -1,392 +0,0 @@
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cpu"
-      prefill_device: "cpu"
-
-# === Rotary Embedding Replacement ===
-
-# GPU 0: layers 0–14
-- match:
-    name: "^model\\.layers\\.([0-9]|1[0-4])\\."
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-
-# GPU 1: layers 15–29
-- match:
-    name: "^model\\.layers\\.(1[5-9]|2[0-9])\\."
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-
-# GPU 2: layers 30–44
-- match:
-    name: "^model\\.layers\\.(3[0-9]|4[0-4])\\."
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
-    kwargs:
-      generate_device: "cuda:2"
-      prefill_device: "cuda:2"
-
-# GPU 3: layers 45–60
-- match:
-    name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\."
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
-    kwargs:
-      generate_device: "cuda:3"
-      prefill_device: "cuda:3"
-
-# === Linear Layers Replacement (excluding self_attn.kv_b_proj) ===
-
-# GPU 0: layers 0–14
-- match:
-    name: "^model\\.layers\\.([0-9]|1[0-4])\\.(?!self_attn\\.kv_b_proj).*$"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-# GPU 1: layers 15–29
-- match:
-    name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.(?!self_attn\\.kv_b_proj).*$"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-# GPU 2: layers 30–44
-- match:
-    name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.(?!self_attn\\.kv_b_proj).*$"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda:2"
-      prefill_device: "cuda:2"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-# GPU 3: layers 45–60
-- match:
-    name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.(?!self_attn\\.kv_b_proj).*$"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda:3"
-      prefill_device: "cuda:3"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-# === MLP (MoE) Replacement ===
-
-# GPU 0: layers 0–14
-- match:
-    name: "^model\\.layers\\.([0-9]|1[0-4])\\.mlp$"
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV3MoE
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-
-# GPU 1: layers 15–29
-- match:
-    name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.mlp$"
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV3MoE
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-
-# GPU 2: layers 30–44
-- match:
-    name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.mlp$"
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV3MoE
-    kwargs:
-      generate_device: "cuda:2"
-      prefill_device: "cuda:2"
-
-# GPU 3: layers 45–60
-- match:
-    name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.mlp$"
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV3MoE
-    kwargs:
-      generate_device: "cuda:3"
-      prefill_device: "cuda:3"
-
-# === MLP Gate Replacement ===
-
-# GPU 0: layers 0–14
-- match:
-    name: "^model\\.layers\\.([0-9]|1[0-4])\\.mlp\\.gate$"
-    class: ktransformers.models.modeling_deepseek_v3.MoEGate
-  replace:
-    class: ktransformers.operators.gate.KMoEGate
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-
-# GPU 1: layers 15–29
-- match:
-    name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.mlp\\.gate$"
-    class: ktransformers.models.modeling_deepseek_v3.MoEGate
-  replace:
-    class: ktransformers.operators.gate.KMoEGate
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-
-# GPU 2: layers 30–44
-- match:
-    name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.mlp\\.gate$"
-    class: ktransformers.models.modeling_deepseek_v3.MoEGate
-  replace:
-    class: ktransformers.operators.gate.KMoEGate
-    kwargs:
-      generate_device: "cuda:2"
-      prefill_device: "cuda:2"
-
-# GPU 3: layers 45–60
-- match:
-    name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.mlp\\.gate$"
-    class: ktransformers.models.modeling_deepseek_v3.MoEGate
-  replace:
-    class: ktransformers.operators.gate.KMoEGate
-    kwargs:
-      generate_device: "cuda:3"
-      prefill_device: "cuda:3"
-
-# === MLP Experts Replacement ===
-# replace with marlin expert. Open and modify layer-num as needed.
-# Each layer of malin experts takes about 6GB of GPU memory.
-# !!!Do remember 'close' cuda graph if you are using marlin expert.!!!
-# !!!KExpertsTorch is untested, we don't have enough VRAM.!!!
-
-# GPU 0: layers 3–4
-# - match:
-#     name: "^model\\.layers\\.([3-4])\\.mlp\\.experts$"
-#   replace:
-#     class: ktransformers.operators.experts.KTransformersExperts
-#     kwargs:
-#       generate_device: "cuda:0"
-#       generate_op:  "KExpertsMarlin"
-#   recursive: False
-
-# # GPU 1: layers 15–17
-# - match:
-#     name: "^model\\.layers\\.(1[5-7])\\.mlp\\.experts$"
-#   replace:
-#     class: ktransformers.operators.experts.KTransformersExperts
-#     kwargs:
-#       generate_device: "cuda:1"
-#       generate_op:  "KExpertsMarlin"
-#   recursive: False
-
-# # GPU 2: layers 30–32
-# - match:
-#     name: "^model\\.layers\\.(3[0-2])\\.mlp\\.experts$"
-#   replace:
-#     class: ktransformers.operators.experts.KTransformersExperts
-#     kwargs:
-#       generate_device: "cuda:2"
-#       generate_op:  "KExpertsMarlin"
-#   recursive: False
-
-# # GPU 3: layers 45–46
-# - match:
-#     name: "^model\\.layers\\.(4[5-6])\\.mlp\\.experts$"
-#   replace:
-#     class: ktransformers.operators.experts.KTransformersExperts
-#     kwargs:
-#       generate_device: "cuda:3"
-#       generate_op:  "KExpertsMarlin"
-#   recursive: False
-
-
-# === MLP Experts Replacement ===
-
-# GPU 0: layers 0–14
-- match:
-    name: "^model\\.layers\\.([0-9]|1[0-4])\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts
-    kwargs:
-      prefill_device: "cuda:0"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KSFTExpertsCPU"
-      out_device: "cuda:0"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False
-
-# GPU 1: layers 15–29
-- match:
-    name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts
-    kwargs:
-      prefill_device: "cuda:1"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KSFTExpertsCPU"
-      out_device: "cuda:1"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False
-
-# GPU 2: layers 30–44
-- match:
-    name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts
-    kwargs:
-      prefill_device: "cuda:2"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KSFTExpertsCPU"
-      out_device: "cuda:2"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False
-
-# GPU 3: layers 45–60
-- match:
-    name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts
-    kwargs:
-      prefill_device: "cuda:3"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KSFTExpertsCPU"
-      out_device: "cuda:3"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False
-
-# === Self-Attention Replacement ===
-
-# GPU 0: layers 0–14
-- match:
-    name: "^model\\.layers\\.([0-9]|1[0-4])\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-      absorb_for_prefill: False
-
-# GPU 1: layers 15–29
-- match:
-    name: "^model\\.layers\\.(1[5-9]|2[0-9])\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-      absorb_for_prefill: False
-
-# GPU 2: layers 30–44
-- match:
-    name: "^model\\.layers\\.(3[0-9]|4[0-4])\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention
-    kwargs:
-      generate_device: "cuda:2"
-      prefill_device: "cuda:2"
-      absorb_for_prefill: False
-
-# GPU 3: layers 45–60
-- match:
-    name: "^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention
-    kwargs:
-      generate_device: "cuda:3"
-      prefill_device: "cuda:3"
-      absorb_for_prefill: False
-
-# === Overall Model Replacement with Transfer Map ===
-
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KDeepseekV2Model"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0 # 0 means close layer‐wise prefill
-      transfer_map:
-        15: "cuda:1" # Layers 15+ on GPU 1
-        30: "cuda:2" # Layers 30+ on GPU 2
-        45: "cuda:3" # Layers 45+ on GPU 3
-
-# === Default Catch-All for Other Modules ===
-
-# GPU 0: layers 0–14
-- match:
-    name: "^model\\.layers\\.([0-9]|1[0-4])\\."
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-
-# GPU 1: layers 15–29
-- match:
-    name: "^model\\.layers\\.(1[5-9]|2[0-9])\\."
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-
-# GPU 2: layers 30–44
-- match:
-    name: "^model\\.layers\\.(3[0-9]|4[0-4])\\."
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cuda:2"
-      prefill_device: "cuda:2"
-
-- match:
-    name: "^lm_head"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda:3"
-      prefill_device: "cuda:3"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-# For final modules (model.norm), ensure they are on GPU 3 (as in your original config)
-- match:
-    name: "(^model\\.layers\\.(4[5-9]|5[0-9]|60)\\.)|(^model\\.norm)"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cuda:3"
-      prefill_device: "cuda:3"

examples/ktransformers/kt_optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml

@@ -1,156 +0,0 @@
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-        generate_device: "cpu"
-        prefill_device: "cpu"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\."
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-- match:
-    name: "^model\\.layers\\.([3456][0-9])\\."
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.(?!self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\.([3456][0-9])\\.(?!self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp$"
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV3MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-- match:
-    name: "^model\\.layers\\.([3456][0-9])\\.mlp$"
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV3MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.gate$"
-    class: ktransformers.models.modeling_deepseek_v3.MoEGate
-  replace:
-    class: ktransformers.operators.gate.KMoEGate
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-- match:
-    name: "^model\\.layers\\.([3456][0-9])\\.mlp\\.gate$"
-    class: ktransformers.models.modeling_deepseek_v3.MoEGate
-  replace:
-    class: ktransformers.operators.gate.KMoEGate     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cuda:0"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op:  "KSFTExpertsCPU"
-      out_device: "cuda:0"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False # don't recursively inject submodules of this module
-
-- match:
-    name: "^model\\.layers\\.([3456][0-9])\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cuda:1"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op:  "KSFTExpertsCPU"
-      out_device: "cuda:1"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False # don't recursively inject submodules of this module
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-- match:
-    name: "^model\\.layers\\.([3456][0-9])\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KDeepseekV2Model"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill
-      transfer_map:
-        30: "cuda:1"
-
-- match:
-    name: "^model\\.layers\\.(0|[1-9]|[12][0-9])\\."
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-
-- match:
-    name: "^lm_head"
-    class: torch.nn.Linear
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "(^model\\.layers\\.([3456][0-9])\\.)|(model.norm)"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cuda:1"
-      prefill_device: "cuda:1"

examples/ktransformers/kt_optimize_rules/DeepSeek-V3-Chat-sft-amx.yaml

@@ -1,77 +0,0 @@
-- match:
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3RotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.YarnRotaryEmbeddingV3
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-
-- match:
-    name: "^lm_head$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-- match:
-    name: "^model\\.layers\\.(?!.*self_attn\\.kv_b_proj).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-- match:
-    name: "^model\\.layers\\..*\\.mlp$"
-    class: ktransformers.models.modeling_deepseek_v3.DeepseekV3MoE
-  replace:
-    class: ktransformers.operators.experts.KDeepseekV3MoE     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    class: ktransformers.models.modeling_deepseek_v3.MoEGate
-  replace:
-    class: ktransformers.operators.gate.KMoEGate
-    kwargs:
-      generate_device: "cuda:0"
-      prefill_device: "cuda:0"
-- match:
-    name: "^model\\.layers\\..*\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cuda"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KSFTExpertsCPU"
-      out_device: "cuda"
-      backend: "AMXInt8" # or "AMXBF16" or "llamafile" (default)
-  recursive: False # don't recursively inject submodules of this module
-- match:
-    name: "^model\\.layers\\..*\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      absorb_for_prefill: False # change this to True to enable long context(prefill may slower).
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KDeepseekV2Model"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cpu"
-      prefill_device: "cpu"

examples/ktransformers/kt_optimize_rules/Qwen3Moe-sft-amx.yaml

@@ -1,80 +0,0 @@
-- match:
-    class: ktransformers.models.modeling_qwen2_moe.Qwen2MoeRotaryEmbedding
-  replace:
-    class: ktransformers.operators.RoPE.RotaryEmbedding
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-
-- match:
-    name: "^lm_head$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-
-# - match:
-#     name: "^model\\.layers\\..*$"  # regular expression
-#     class: torch.nn.Linear  # only match modules matching name and class simultaneously
-#   replace:
-#     class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-#     kwargs:
-#       generate_device: "cuda"
-#       prefill_device: "cuda"
-#       generate_op: "KLinearTorch"
-#       prefill_op: "KLinearTorch"
-- match:
-    name: "^model\\.layers\\.(?!.*mlp\\.shared_expert_gate).*$"  # regular expression
-    class: torch.nn.Linear  # only match modules matching name and class simultaneously
-  replace:
-    class: ktransformers.operators.linear.KTransformersLinear  # optimized Kernel on quantized data types
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-      generate_op: "KLinearTorch"
-      prefill_op: "KLinearTorch"
-- match:
-    name: "^model\\.layers\\..*\\.mlp$"
-  replace:
-    class: ktransformers.operators.experts.KQwen3MoeSparseMoeBlock     # mlp module with custom forward function
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-
-- match:
-    name: "^model\\.layers\\..*\\.mlp\\.experts$"
-  replace:
-    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
-    kwargs:
-      prefill_device: "cuda"
-      prefill_op: "KExpertsTorch"
-      generate_device: "cpu"
-      generate_op: "KSFTExpertsCPU"
-      out_device: "cuda"
-      backend: "AMXInt8" # or "AMXBF16" or "AMXInt8"
-  recursive: False # don't recursively inject submodules of this module
-- match:
-    name: "^model\\.layers\\..*\\.self_attn$"
-  replace:
-    class: ktransformers.operators.attention.KQwen3MoeAttention # optimized MLA implementation
-    kwargs:
-      generate_device: "cuda"
-      prefill_device: "cuda"
-- match:
-    name: "^model.embed_tokens"
-  replace:
-    class: "default"
-    kwargs:
-      generate_device: "cpu"
-      prefill_device: "cpu"
-
-- match:
-    name: "^model$"
-  replace:
-    class: "ktransformers.operators.models.KQwen3MoeModel"
-    kwargs:
-      per_layer_prefill_intput_threshold: 0

examples/ktransformers/train_lora/deepseek_v2_lora_sft_kt.yaml

@@ -19,7 +19,7 @@ preprocessing_num_workers: 16
 dataloader_num_workers: 4
 
 ### output
-output_dir: saves/Kllama_deepseekV2
+output_dir: saves/KT_FT_deepseekV2
 logging_steps: 10
 save_steps: 500
 plot_loss: true
@@ -39,14 +39,7 @@ ddp_timeout: 180000000
 resume_from_checkpoint: null
 
 ### ktransformers
-use_kt: true # use KTransformers as LoRA sft backend
-kt_optimize_rule: examples/ktransformers/kt_optimize_rules/DeepSeek-V2-Lite-Chat-sft-amx.yaml
-cpu_infer: 32
-chunk_size: 8192
-
-### eval
-# eval_dataset: alpaca_en_demo
-# val_size: 0.1
-# per_device_eval_batch_size: 1
-# eval_strategy: steps
-# eval_steps: 500
+use_kt: true
+# Pair with fsdp2_kt_bf16.yaml for original BF16 checkpoints.
+# For pre-converted expert weights, uncomment kt_weight_path and use fsdp2_kt_int8.yaml or fsdp2_kt_int4.yaml.
+# kt_weight_path: /path/to/DeepSeek-V2-Lite-AMXINT8

examples/ktransformers/train_lora/deepseek_v3_lora_sft_kt.yaml

@@ -1,5 +1,5 @@
 ### model
-model_name_or_path: opensourcerelease/DeepSeek-V3-bf16
+model_name_or_path: deepseek-ai/DeepSeek-V3-0324-BF16 # need to convert to BF16 checkpoint first
 trust_remote_code: true
 
 ### method
@@ -19,7 +19,7 @@ preprocessing_num_workers: 16
 dataloader_num_workers: 4
 
 ### output
-output_dir: saves/Kllama_deepseekV3
+output_dir: saves/KT_FT_deepseekV3
 logging_steps: 10
 save_steps: 500
 plot_loss: true
@@ -39,14 +39,7 @@ ddp_timeout: 180000000
 resume_from_checkpoint: null
 
 ### ktransformers
-use_kt: true # use KTransformers as LoRA sft backend
-kt_optimize_rule: examples/ktransformers/kt_optimize_rules/DeepSeek-V3-Chat-sft-amx-multi-gpu.yaml
-cpu_infer: 32
-chunk_size: 8192
-
-### eval
-# eval_dataset: alpaca_en_demo
-# val_size: 0.1
-# per_device_eval_batch_size: 1
-# eval_strategy: steps
-# eval_steps: 500
+use_kt: true
+# Pair with fsdp2_kt_bf16.yaml for original BF16 checkpoints.
+# For pre-converted expert weights, uncomment kt_weight_path and use fsdp2_kt_int8.yaml or fsdp2_kt_int4.yaml.
+# kt_weight_path: /path/to/DeepSeek-V3-AMXINT8

examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml

@@ -0,0 +1,46 @@
+### model
+model_name_or_path: Qwen/Qwen3.5-397B-A17B
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: lora
+lora_rank: 8
+lora_target: all
+
+### dataset
+dataset: identity, alpaca_en_demo
+template: qwen3_5
+cutoff_len: 2048
+max_samples: 100000
+overwrite_cache: true
+preprocessing_num_workers: 16
+dataloader_num_workers: 4
+
+### output
+output_dir: saves/KT_FT_qwen35Moe
+logging_steps: 10
+save_steps: 500
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 8
+learning_rate: 1.0e-4
+num_train_epochs: 3.0
+lr_scheduler_type: cosine
+warmup_ratio: 0.1
+bf16: true
+ddp_timeout: 180000000
+resume_from_checkpoint: null
+
+### ktransformers
+use_kt: true
+# For original BF16 checkpoints, start with examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml.
+# For pre-converted expert weights, uncomment kt_weight_path and use fsdp2_kt_int8.yaml or fsdp2_kt_int4.yaml.
+# Pair the 397B path with fsdp2_kt_int8.yaml, tune cutoff_len to prepared weights and GPU memory.
+# kt_weight_path: /path/to/Qwen3.5-MoE-AMXINT8

examples/ktransformers/train_lora/qwen3moe_lora_sft_kt.yaml

@@ -11,7 +11,7 @@ lora_target: all
 
 ### dataset
 dataset: identity, alpaca_en_demo
-template: qwen3_nothink
+template: qwen3
 cutoff_len: 2048
 max_samples: 100000
 overwrite_cache: true
@@ -19,9 +19,9 @@ preprocessing_num_workers: 16
 dataloader_num_workers: 4
 
 ### output
-output_dir: saves/Kllama_Qwen3MoE_235bA22b
+output_dir: saves/KT_FT_qwen3Moe
 logging_steps: 10
-save_steps: 200
+save_steps: 500
 plot_loss: true
 overwrite_output_dir: true
 save_only_model: false
@@ -31,7 +31,7 @@ report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
 per_device_train_batch_size: 1
 gradient_accumulation_steps: 8
 learning_rate: 1.0e-4
-num_train_epochs: 3
+num_train_epochs: 3.0
 lr_scheduler_type: cosine
 warmup_ratio: 0.1
 bf16: true
@@ -39,14 +39,7 @@ ddp_timeout: 180000000
 resume_from_checkpoint: null
 
 ### ktransformers
-use_kt: true # use KTransformers as LoRA sft backend
-kt_optimize_rule: examples/ktransformers/kt_optimize_rules/Qwen3Moe-sft-amx.yaml
-cpu_infer: 32
-chunk_size: 8192
-
-### eval
-# eval_dataset: alpaca_en_demo
-# val_size: 0.1
-# per_device_eval_batch_size: 1
-# eval_strategy: steps
-# eval_steps: 500
+use_kt: true
+# Pair with examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml for original BF16 checkpoints.
+# For pre-converted expert weights, uncomment kt_weight_path and use fsdp2_kt_int8.yaml or fsdp2_kt_int4.yaml.
+# kt_weight_path: /path/to/Qwen3-235B-A22B-Instruct-2507-AMXINT8