[GH-ISSUE #76] 8-GPU configuration on L40 OOM #38

New issue

Closed

opened 2026-10-29 14:32:54 +00:00 by bsakaguchi · 8 comments

bsakaguchi commented 2026-10-29 14:32:54 +00:00

Owner

Copy link

Originally created by @fengyang95 on GitHub (Sep 4, 2024).
Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/76

I tried running deepseek-v2 with an 8xL40 46G configuration, but I encountered a GPU memory Out-of-Memory (OOM) error. Why would such a large amount of GPU memory still lead to an OOM issue?

- match:
    name: "^model.embed_tokens"
  replace:
    class: "default"
    kwargs:
        generate_device: "cuda:0"
        prefill_device: "cuda:0"

- match:
    name: "^model\\.layers\\.([0-3])\\."
    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\\.([4-9]|[1][0-1])\\."
    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\\.([1][2-9])\\."
    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
  replace:
    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
    kwargs:
      generate_device: "cuda:2"
      prefill_device: "cuda:2"
- match:
    name: "^model\\.layers\\.([2][0-7])\\."
    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
  replace:
    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
    kwargs:
      generate_device: "cuda:3"
      prefill_device: "cuda:3"
- match:
    name: "^model\\.layers\\.([2][8-9]|[3][0-5])\\."
    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
  replace:
    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
    kwargs:
      generate_device: "cuda:4"
      prefill_device: "cuda:4"
- match:
    name: "^model\\.layers\\.([3][6-9]|[4][0-3])\\."
    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
  replace:
    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
    kwargs:
      generate_device: "cuda:5"
      prefill_device: "cuda:5"
- match:
    name: "^model\\.layers\\.([4][4-9]|[5][0-1])\\."
    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
  replace:
    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
    kwargs:
      generate_device: "cuda:6"
      prefill_device: "cuda:6"
- match:
    name: "^model\\.layers\\.([5][2-9])\\."
    class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding
  replace:
    class: ktransformers.operators.RoPE.YarnRotaryEmbedding
    kwargs:
      generate_device: "cuda:7"
      prefill_device: "cuda:7"


- match:
    name: "^model\\.layers\\.([0-3])\\.(?!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: "KLinearMarlin"
      prefill_op: "KLinearTorch"
- match:
    name: "^model\\.layers\\.([4-9]|[1][0-1])\\.(?!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: "KLinearMarlin"
      prefill_op: "KLinearTorch"
- match:
    name: "^model\\.layers\\.([1][2-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:2"
      prefill_device: "cuda:2"
      generate_op: "KLinearMarlin"
      prefill_op: "KLinearTorch"
- match:
    name: "^model\\.layers\\.([2][0-7])\\.(?!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:3"
      prefill_device: "cuda:3"
      generate_op: "KLinearMarlin"
      prefill_op: "KLinearTorch"
- match:
    name: "^model\\.layers\\.([2][8-9]|[3][0-5])\\.(?!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:4"
      prefill_device: "cuda:4"
      generate_op: "KLinearMarlin"
      prefill_op: "KLinearTorch"
- match:
    name: "^model\\.layers\\.([3][6-9]|[4][0-3])\\.(?!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:5"
      prefill_device: "cuda:5"
      generate_op: "KLinearMarlin"
      prefill_op: "KLinearTorch"
- match:
    name: "^model\\.layers\\.([4][4-9]|[5][0-1])\\.(?!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:6"
      prefill_device: "cuda:6"
      generate_op: "KLinearMarlin"
      prefill_op: "KLinearTorch"
- match:
    name: "^model\\.layers\\.([5][2-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:7"
      prefill_device: "cuda:7"
      generate_op: "KLinearMarlin"
      prefill_op: "KLinearTorch"

- match:
    name: "^model\\.layers\\.([0-3])\\.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\\.([4-9]|[1][0-1])\\.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\\.([1][2-9])\\.mlp$"
    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
  replace:
    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
    kwargs:
      generate_device: "cuda:2"
      prefill_device: "cuda:2"
- match:
    name: "^model\\.layers\\.([2][0-7])\\.mlp$"
    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
  replace:
    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
    kwargs:
      generate_device: "cuda:3"
      prefill_device: "cuda:3"

- match:
    name: "^model\\.layers\\.([2][8-9]|[3][0-5])\\.mlp$"
    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
  replace:
    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
    kwargs:
      generate_device: "cuda:4"
      prefill_device: "cuda:4"
- match:
    name: "^model\\.layers\\.([3][6-9]|[4][0-3])\\.mlp$"
    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
  replace:
    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
    kwargs:
      generate_device: "cuda:5"
      prefill_device: "cuda:5"
- match:
    name: "^model\\.layers\\.([4][4-9]|[5][0-1])\\.mlp$"
    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
  replace:
    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
    kwargs:
      generate_device: "cuda:6"
      prefill_device: "cuda:6"
- match:
    name: "^model\\.layers\\.([5][2-9])\\.mlp$"
    class: ktransformers.models.modeling_deepseek.DeepseekV2MoE
  replace:
    class: ktransformers.operators.experts.KDeepseekV2MoE     # mlp module with custom forward function
    kwargs:
      generate_device: "cuda:7"
      prefill_device: "cuda:7"

- match:
    name: "^model\\.layers\\.([0-3])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
    kwargs:
      prefill_device: "cuda:0"
      prefill_op: "KExpertsTorch"
      generate_device: "cuda:0"
      generate_op:  "KExpertsTorch"
      out_device: "cuda:0"
  recursive: False # don't recursively inject submodules of this module
- match:
    name: "^model\\.layers\\.([4-9]|[1][0-1])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
    kwargs:
      prefill_device: "cuda:1"
      prefill_op: "KExpertsTorch"
      generate_device: "cuda:1"
      generate_op:  "KExpertsTorch"
      out_device: "cuda:1"
  recursive: False # don't recursively inject submodules of this module
- match:
    name: "^model\\.layers\\.([1][2-9])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
    kwargs:
      prefill_device: "cuda:2"
      prefill_op: "KExpertsTorch"
      generate_device: "cuda:2"
      generate_op:  "KExpertsTorch"
      out_device: "cuda:2"
  recursive: False # don't recursively inject submodules of this module
- match:
    name: "^model\\.layers\\.([2][0-7])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
    kwargs:
      prefill_device: "cuda:3"
      prefill_op: "KExpertsTorch"
      generate_device: "cuda:3"
      generate_op:  "KExpertsTorch"
      out_device: "cuda:3"
  recursive: False # don't recursively inject submodules of this module
- match:
    name: "^model\\.layers\\.([2][8-9]|[3][0-5])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
    kwargs:
      prefill_device: "cuda:4"
      prefill_op: "KExpertsTorch"
      generate_device: "cuda:4"
      generate_op:  "KExpertsTorch"
      out_device: "cuda:4"
  recursive: False # don't recursively inject submodules of this module
- match:
    name: "^model\\.layers\\.([3][6-9]|[4][0-3])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
    kwargs:
      prefill_device: "cuda:5"
      prefill_op: "KExpertsTorch"
      generate_device: "cuda:5"
      generate_op:  "KExpertsTorch"
      out_device: "cuda:5"
  recursive: False # don't recursively inject submodules of this module
- match:
    name: "^model\\.layers\\.([4][4-9]|[5][0-1])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
    kwargs:
      prefill_device: "cuda:6"
      prefill_op: "KExpertsTorch"
      generate_device: "cuda:6"
      generate_op:  "KExpertsTorch"
      out_device: "cuda:6"
  recursive: False # don't recursively inject submodules of this module
- match:
    name: "^model\\.layers\\.([5][2-9])\\.mlp\\.experts$"
  replace:
    class: ktransformers.operators.experts.KTransformersExperts     # custom MoE Kernel with expert paralleism
    kwargs:
      prefill_device: "cuda:7"
      prefill_op: "KExpertsTorch"
      generate_device: "cuda:7"
      generate_op:  "KExpertsTorch"
      out_device: "cuda:7"
  recursive: False # don't recursively inject submodules of this module

- match:
    name: "^model\\.layers\\.([0-3])\\.self_attn$"
  replace:
    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
    kwargs:
      generate_device: "cuda:0"
      prefill_device: "cuda:0"
- match:
    name: "^model\\.layers\\.([4-9]|[1][0-1])\\.self_attn$"
  replace:
    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
    kwargs:
      generate_device: "cuda:1"
      prefill_device: "cuda:1"
- match:
    name: "^model\\.layers\\.([1][2-9])\\.self_attn$"
  replace:
    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
    kwargs:
      generate_device: "cuda:2"
      prefill_device: "cuda:2"
- match:
    name: "^model\\.layers\\.([2][0-7])\\.self_attn$"
  replace:
    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
    kwargs:
      generate_device: "cuda:3"
      prefill_device: "cuda:3"

- match:
    name: "^model\\.layers\\.([2][8-9]|[3][0-5])\\.self_attn$"
  replace:
    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
    kwargs:
      generate_device: "cuda:4"
      prefill_device: "cuda:4"
- match:
    name: "^model\\.layers\\.([3][6-9]|[4][0-3])\\.self_attn$"
  replace:
    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
    kwargs:
      generate_device: "cuda:5"
      prefill_device: "cuda:5"
- match:
    name: "^model\\.layers\\.([4][4-9]|[5][0-1])\\.self_attn$"
  replace:
    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
    kwargs:
      generate_device: "cuda:6"
      prefill_device: "cuda:6"
- match:
    name: "^model\\.layers\\.([5][2-9])\\.self_attn$"
  replace:
    class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation
    kwargs:
      generate_device: "cuda:7"
      prefill_device: "cuda:7"


- match:
    name: "^model$"
  replace:
    class: "ktransformers.operators.models.KDeepseekV2Model"
    kwargs:
      per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill
      transfer_map: 
        4: "cuda:1"
        12: "cuda:2"
        20: "cuda:3"
        28: "cuda:4"
        36: "cuda:5"
        44: "cuda:6"
        52: "cuda:7"

- match:
    name: "^model\\.layers\\.([0-3])\\."
  replace:
    class: "default"
    kwargs:
      generate_device: "cuda:0"
      prefill_device: "cuda:0"
- match:
    name: "(^model\\.layers\\.([4-9]|[1][0-1])\\.)"
  replace:
    class: "default"
    kwargs:
      generate_device: "cuda:1"
      prefill_device: "cuda:1"
- match:
    name: "(^model\\.layers\\.([1][2-9])\\.)"
  replace:
    class: "default"
    kwargs:
      generate_device: "cuda:2"
      prefill_device: "cuda:2"
- match:
    name: "(^model\\.layers\\.([2][0-7])\\.)"
  replace:
    class: "default"
    kwargs:
      generate_device: "cuda:3"
      prefill_device: "cuda:3"

- match:
    name: "(^model\\.layers\\.([2][8-9]|[3][0-5])\\.)"
  replace:
    class: "default"
    kwargs:
      generate_device: "cuda:4"
      prefill_device: "cuda:4"
- match:
    name: "(^model\\.layers\\.([3][6-9]|[4][0-3])\\.)"
  replace:
    class: "default"
    kwargs:
      generate_device: "cuda:5"
      prefill_device: "cuda:5"
- match:
    name: "(^model\\.layers\\.([4][4-9]|[5][0-1])\\.)"
  replace:
    class: "default"
    kwargs:
      generate_device: "cuda:6"
      prefill_device: "cuda:6"
- match:
    name: "(^model\\.layers\\.([5][2-9])\\.)|(^model.norm)|(^lm_head)"
  replace:
    class: "default"
    kwargs:
      generate_device: "cuda:7"
      prefill_device: "cuda:7"

Originally created by @fengyang95 on GitHub (Sep 4, 2024). Original GitHub issue: https://github.com/kvcache-ai/ktransformers/issues/76 I tried running deepseek-v2 with an 8xL40 46G configuration, but I encountered a GPU memory Out-of-Memory (OOM) error. Why would such a large amount of GPU memory still lead to an OOM issue? ``` yaml - match: name: "^model.embed_tokens" replace: class: "default" kwargs: generate_device: "cuda:0" prefill_device: "cuda:0" - match: name: "^model\\.layers\\.([0-3])\\." 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\\.([4-9]|[1][0-1])\\." 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\\.([1][2-9])\\." class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbedding kwargs: generate_device: "cuda:2" prefill_device: "cuda:2" - match: name: "^model\\.layers\\.([2][0-7])\\." class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbedding kwargs: generate_device: "cuda:3" prefill_device: "cuda:3" - match: name: "^model\\.layers\\.([2][8-9]|[3][0-5])\\." class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbedding kwargs: generate_device: "cuda:4" prefill_device: "cuda:4" - match: name: "^model\\.layers\\.([3][6-9]|[4][0-3])\\." class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbedding kwargs: generate_device: "cuda:5" prefill_device: "cuda:5" - match: name: "^model\\.layers\\.([4][4-9]|[5][0-1])\\." class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbedding kwargs: generate_device: "cuda:6" prefill_device: "cuda:6" - match: name: "^model\\.layers\\.([5][2-9])\\." class: ktransformers.models.modeling_deepseek.DeepseekV2YarnRotaryEmbedding replace: class: ktransformers.operators.RoPE.YarnRotaryEmbedding kwargs: generate_device: "cuda:7" prefill_device: "cuda:7" - match: name: "^model\\.layers\\.([0-3])\\.(?!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: "KLinearMarlin" prefill_op: "KLinearTorch" - match: name: "^model\\.layers\\.([4-9]|[1][0-1])\\.(?!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: "KLinearMarlin" prefill_op: "KLinearTorch" - match: name: "^model\\.layers\\.([1][2-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:2" prefill_device: "cuda:2" generate_op: "KLinearMarlin" prefill_op: "KLinearTorch" - match: name: "^model\\.layers\\.([2][0-7])\\.(?!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:3" prefill_device: "cuda:3" generate_op: "KLinearMarlin" prefill_op: "KLinearTorch" - match: name: "^model\\.layers\\.([2][8-9]|[3][0-5])\\.(?!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:4" prefill_device: "cuda:4" generate_op: "KLinearMarlin" prefill_op: "KLinearTorch" - match: name: "^model\\.layers\\.([3][6-9]|[4][0-3])\\.(?!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:5" prefill_device: "cuda:5" generate_op: "KLinearMarlin" prefill_op: "KLinearTorch" - match: name: "^model\\.layers\\.([4][4-9]|[5][0-1])\\.(?!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:6" prefill_device: "cuda:6" generate_op: "KLinearMarlin" prefill_op: "KLinearTorch" - match: name: "^model\\.layers\\.([5][2-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:7" prefill_device: "cuda:7" generate_op: "KLinearMarlin" prefill_op: "KLinearTorch" - match: name: "^model\\.layers\\.([0-3])\\.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\\.([4-9]|[1][0-1])\\.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\\.([1][2-9])\\.mlp$" class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function kwargs: generate_device: "cuda:2" prefill_device: "cuda:2" - match: name: "^model\\.layers\\.([2][0-7])\\.mlp$" class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function kwargs: generate_device: "cuda:3" prefill_device: "cuda:3" - match: name: "^model\\.layers\\.([2][8-9]|[3][0-5])\\.mlp$" class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function kwargs: generate_device: "cuda:4" prefill_device: "cuda:4" - match: name: "^model\\.layers\\.([3][6-9]|[4][0-3])\\.mlp$" class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function kwargs: generate_device: "cuda:5" prefill_device: "cuda:5" - match: name: "^model\\.layers\\.([4][4-9]|[5][0-1])\\.mlp$" class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function kwargs: generate_device: "cuda:6" prefill_device: "cuda:6" - match: name: "^model\\.layers\\.([5][2-9])\\.mlp$" class: ktransformers.models.modeling_deepseek.DeepseekV2MoE replace: class: ktransformers.operators.experts.KDeepseekV2MoE # mlp module with custom forward function kwargs: generate_device: "cuda:7" prefill_device: "cuda:7" - match: name: "^model\\.layers\\.([0-3])\\.mlp\\.experts$" replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: prefill_device: "cuda:0" prefill_op: "KExpertsTorch" generate_device: "cuda:0" generate_op: "KExpertsTorch" out_device: "cuda:0" recursive: False # don't recursively inject submodules of this module - match: name: "^model\\.layers\\.([4-9]|[1][0-1])\\.mlp\\.experts$" replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: prefill_device: "cuda:1" prefill_op: "KExpertsTorch" generate_device: "cuda:1" generate_op: "KExpertsTorch" out_device: "cuda:1" recursive: False # don't recursively inject submodules of this module - match: name: "^model\\.layers\\.([1][2-9])\\.mlp\\.experts$" replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: prefill_device: "cuda:2" prefill_op: "KExpertsTorch" generate_device: "cuda:2" generate_op: "KExpertsTorch" out_device: "cuda:2" recursive: False # don't recursively inject submodules of this module - match: name: "^model\\.layers\\.([2][0-7])\\.mlp\\.experts$" replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: prefill_device: "cuda:3" prefill_op: "KExpertsTorch" generate_device: "cuda:3" generate_op: "KExpertsTorch" out_device: "cuda:3" recursive: False # don't recursively inject submodules of this module - match: name: "^model\\.layers\\.([2][8-9]|[3][0-5])\\.mlp\\.experts$" replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: prefill_device: "cuda:4" prefill_op: "KExpertsTorch" generate_device: "cuda:4" generate_op: "KExpertsTorch" out_device: "cuda:4" recursive: False # don't recursively inject submodules of this module - match: name: "^model\\.layers\\.([3][6-9]|[4][0-3])\\.mlp\\.experts$" replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: prefill_device: "cuda:5" prefill_op: "KExpertsTorch" generate_device: "cuda:5" generate_op: "KExpertsTorch" out_device: "cuda:5" recursive: False # don't recursively inject submodules of this module - match: name: "^model\\.layers\\.([4][4-9]|[5][0-1])\\.mlp\\.experts$" replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: prefill_device: "cuda:6" prefill_op: "KExpertsTorch" generate_device: "cuda:6" generate_op: "KExpertsTorch" out_device: "cuda:6" recursive: False # don't recursively inject submodules of this module - match: name: "^model\\.layers\\.([5][2-9])\\.mlp\\.experts$" replace: class: ktransformers.operators.experts.KTransformersExperts # custom MoE Kernel with expert paralleism kwargs: prefill_device: "cuda:7" prefill_op: "KExpertsTorch" generate_device: "cuda:7" generate_op: "KExpertsTorch" out_device: "cuda:7" recursive: False # don't recursively inject submodules of this module - match: name: "^model\\.layers\\.([0-3])\\.self_attn$" replace: class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation kwargs: generate_device: "cuda:0" prefill_device: "cuda:0" - match: name: "^model\\.layers\\.([4-9]|[1][0-1])\\.self_attn$" replace: class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation kwargs: generate_device: "cuda:1" prefill_device: "cuda:1" - match: name: "^model\\.layers\\.([1][2-9])\\.self_attn$" replace: class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation kwargs: generate_device: "cuda:2" prefill_device: "cuda:2" - match: name: "^model\\.layers\\.([2][0-7])\\.self_attn$" replace: class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation kwargs: generate_device: "cuda:3" prefill_device: "cuda:3" - match: name: "^model\\.layers\\.([2][8-9]|[3][0-5])\\.self_attn$" replace: class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation kwargs: generate_device: "cuda:4" prefill_device: "cuda:4" - match: name: "^model\\.layers\\.([3][6-9]|[4][0-3])\\.self_attn$" replace: class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation kwargs: generate_device: "cuda:5" prefill_device: "cuda:5" - match: name: "^model\\.layers\\.([4][4-9]|[5][0-1])\\.self_attn$" replace: class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation kwargs: generate_device: "cuda:6" prefill_device: "cuda:6" - match: name: "^model\\.layers\\.([5][2-9])\\.self_attn$" replace: class: ktransformers.operators.attention.KDeepseekV2Attention # optimized MLA implementation kwargs: generate_device: "cuda:7" prefill_device: "cuda:7" - match: name: "^model$" replace: class: "ktransformers.operators.models.KDeepseekV2Model" kwargs: per_layer_prefill_intput_threshold: 0 # 0 is close layer wise prefill transfer_map: 4: "cuda:1" 12: "cuda:2" 20: "cuda:3" 28: "cuda:4" 36: "cuda:5" 44: "cuda:6" 52: "cuda:7" - match: name: "^model\\.layers\\.([0-3])\\." replace: class: "default" kwargs: generate_device: "cuda:0" prefill_device: "cuda:0" - match: name: "(^model\\.layers\\.([4-9]|[1][0-1])\\.)" replace: class: "default" kwargs: generate_device: "cuda:1" prefill_device: "cuda:1" - match: name: "(^model\\.layers\\.([1][2-9])\\.)" replace: class: "default" kwargs: generate_device: "cuda:2" prefill_device: "cuda:2" - match: name: "(^model\\.layers\\.([2][0-7])\\.)" replace: class: "default" kwargs: generate_device: "cuda:3" prefill_device: "cuda:3" - match: name: "(^model\\.layers\\.([2][8-9]|[3][0-5])\\.)" replace: class: "default" kwargs: generate_device: "cuda:4" prefill_device: "cuda:4" - match: name: "(^model\\.layers\\.([3][6-9]|[4][0-3])\\.)" replace: class: "default" kwargs: generate_device: "cuda:5" prefill_device: "cuda:5" - match: name: "(^model\\.layers\\.([4][4-9]|[5][0-1])\\.)" replace: class: "default" kwargs: generate_device: "cuda:6" prefill_device: "cuda:6" - match: name: "(^model\\.layers\\.([5][2-9])\\.)|(^model.norm)|(^lm_head)" replace: class: "default" kwargs: generate_device: "cuda:7" prefill_device: "cuda:7" ```

bsakaguchi closed this issue 2026-10-29 14:32:54 +00:00

bsakaguchi commented 2026-10-29 14:32:55 +00:00

Author

Owner

Copy link

@Azure-Tang commented on GitHub (Sep 4, 2024):

Ok, if you want to put all modules to gpu, then your required VRAM per layer is approximately 5.2G. So 9 layers on one GPU may cause oom. You can offload one layer of experts to CPU

<!-- gh-comment-id:2329226965 --> @Azure-Tang commented on GitHub (Sep 4, 2024): Ok, if you want to put all modules to gpu, then your required VRAM per layer is approximately 5.2G. So 9 layers on one GPU may cause oom. You can offload one layer of experts to CPU

bsakaguchi commented 2026-10-29 14:32:55 +00:00

Author

Owner

Copy link

@fengyang95 commented on GitHub (Sep 5, 2024):

Ok, if you want to put all modules to gpu, then your required VRAM per layer is approximately 5.2G. So 9 layers on one GPU may cause oom. You can offload one layer of experts to CPU
@Azure-Tang But I only put 8 layers on each GPU, why is it still OOM? Do I need to reserve some space?

<!-- gh-comment-id:2330400029 --> @fengyang95 commented on GitHub (Sep 5, 2024): > Ok, if you want to put all modules to gpu, then your required VRAM per layer is approximately 5.2G. So 9 layers on one GPU may cause oom. You can offload one layer of experts to CPU @Azure-Tang But I only put 8 layers on each GPU, why is it still OOM? Do I need to reserve some space?

bsakaguchi commented 2026-10-29 14:32:55 +00:00

Author

Owner

Copy link

@Azure-Tang commented on GitHub (Sep 5, 2024):

The 5.2G figure is an approximation, and the actual value may be slightly higher. As a result, 8 layers per GPU could be the upper limit. You can consider offloading one layer’s experts to the CPU for each GPU. This way, each CPU will handle 8 layers of experts modules, with the remaining layers calculated in the GPU.

<!-- gh-comment-id:2330459977 --> @Azure-Tang commented on GitHub (Sep 5, 2024): The 5.2G figure is an approximation, and the actual value may be slightly higher. As a result, 8 layers per GPU could be the upper limit. You can consider offloading one layer’s experts to the CPU for each GPU. This way, each CPU will handle 8 layers of experts modules, with the remaining layers calculated in the GPU.

bsakaguchi commented 2026-10-29 14:32:55 +00:00

Author

Owner

Copy link

@fengyang95 commented on GitHub (Sep 5, 2024):

The 5.2G figure is an approximation, and the actual value may be slightly higher. As a result, 8 layers per GPU could be the upper limit. You can consider offloading one layer’s experts to the CPU for each GPU. This way, each CPU will handle 8 layers of experts modules, with the remaining layers calculated in the GPU.

Do you mean to allocate 7 layers to each GPU and offload 1 layer to the CPU?

<!-- gh-comment-id:2330468541 --> @fengyang95 commented on GitHub (Sep 5, 2024): > The 5.2G figure is an approximation, and the actual value may be slightly higher. As a result, 8 layers per GPU could be the upper limit. You can consider offloading one layer’s experts to the CPU for each GPU. This way, each CPU will handle 8 layers of experts modules, with the remaining layers calculated in the GPU. Do you mean to allocate 7 layers to each GPU and offload 1 layer to the CPU?

bsakaguchi commented 2026-10-29 14:32:55 +00:00

Author

Owner

Copy link

@Azure-Tang commented on GitHub (Sep 5, 2024):

Not a whole layer to the CPU, only offload experts module in one decoder layer for each 8-layer GPU seems enough.

<!-- gh-comment-id:2330493255 --> @Azure-Tang commented on GitHub (Sep 5, 2024): Not `a whole layer` to the CPU, only offload `experts` module in one decoder layer for each 8-layer GPU seems enough.

bsakaguchi commented 2026-10-29 14:32:56 +00:00

Author

Owner

Copy link

@fengyang95 commented on GitHub (Sep 5, 2024):

Not a whole layer to the CPU, only offload experts module in one decoder layer for each 8-layer GPU seems enough.

I didn't quite understand your point. Could you pls demonstrate it using a configuration?

<!-- gh-comment-id:2331357109 --> @fengyang95 commented on GitHub (Sep 5, 2024): > Not `a whole layer` to the CPU, only offload `experts` module in one decoder layer for each 8-layer GPU seems enough. I didn't quite understand your point. Could you pls demonstrate it using a configuration?

bsakaguchi commented 2026-10-29 14:32:56 +00:00

Author

Owner

Copy link

@Azure-Tang commented on GitHub (Sep 6, 2024):

For example, you got 8 layers for each GPU, which cause oom. And your 8 layers' parameter distribution is like this:

Then you offload one layer's experts module to cpu (which we have done in our example yamls), so the layer will be like this：

And keep other 7 whole layers on GPU, which like this:

Which will offload some parameters to you ram. If it still oom, you can do this for another layer.

<!-- gh-comment-id:2333066103 --> @Azure-Tang commented on GitHub (Sep 6, 2024): For example, you got 8 layers for each GPU, which cause oom. And your 8 layers' parameter distribution is like this: <img width="325" alt="image" src="https://github.com/user-attachments/assets/c931a2fa-41fb-4e6d-9027-2223d30bf4ee"> Then you offload one layer's experts module to cpu (which we have done in our example yamls), so the layer will be like this： <img width="286" alt="image" src="https://github.com/user-attachments/assets/f4581b5e-0abb-4c35-8c81-d4eaded4bc4d"> And keep other 7 whole layers on GPU, which like this: <img width="325" alt="image" src="https://github.com/user-attachments/assets/c931a2fa-41fb-4e6d-9027-2223d30bf4ee"> Which will offload some parameters to you ram. If it still oom, you can do this for another layer.

bsakaguchi commented 2026-10-29 14:32:56 +00:00

Author

Owner

Copy link

@Azure-Tang commented on GitHub (Sep 11, 2024):

Hi, I’ll be closing this topic as there hasn’t been any response～

<!-- gh-comment-id:2342770960 --> @Azure-Tang commented on GitHub (Sep 11, 2024): Hi, I’ll be closing this topic as there hasn’t been any response～

bsakaguchi referenced this issue 2026-10-29 14:51:27 +00:00

[PR #38] [MERGED] [update] README #1327

Sign in to join this conversation.

No Branch/Tag specified

Branches Tags

main

auto/sync-sglang

gh-pages

docs/kt-ft-figure2-single-column-20261010

perf/bf16-grad-norm-20261010

docs/kt-ft-paper-20261008

release/20260911-carrier-repair

release/20260911-ci-python

release/20260911-pypi

ci/four-main-pypi-release

ci/model-e2e-acceptance

release/0.7.0.post3-frozen-pypi

fix/glm53-index-share-sync

release/0.7.0.post2-frozen-pypi

hotfix/0.7.0.post2-pypi-runtime

fix/avx2-wheel-feature-isolation

fix/release-self-hosted-runner

release/automate-kt-sft-pypi

support-glm5-next-djw

docs/deepseek-v4-flash-docker

ovowei-patch-3

support-glm52-djw

doc-reorg-entrypoints-support-matrix

bump-sglang-pr31

docs-v061-refresh

fp4-moe-amx

flatten-ktransformers-shim

archive-kt-sft

kt-kernel-torch-range-v061

sft-whl

sft

docs/gosim-2026-and-roadmap-link

ovowei-patch-2

ovowei-patch-1

revert-1918-fix/rocm-cpuinfer-stream-bridge

chore/kt-layerwise-prefill-main

chore/kt-layerwise-prefill-label

sft-rel

support-glm51-djw

release/v0.5.3

chore/kt-run-numa-nodes-unified

mtp

ksft-sglang

avx2

fix/sglang-kt-detection

layerwise-prefill-opt

feat/mistral-loader-support

fix-docker-build

minimax-m2.5

qwen3.5

add-maintainers

oql-doc

cli-dev-ljq

ljq-cli-dev-new

JimmyPeilinLi-patch-2

kimi_k2.5

oql-dev

experts_sched

JimmyPeilinLi-patch-1

ErvinXie-patch-1

ksft-llama

glm47_fp8

debug_llama

refactor_sft

ljq-cli-dev

refactor_sft_int4_debug

lut-nvfp4

fix-import

fix-deploy-djw

cpu-fix

ljq-dev-1

route-moe

sosp25-ae

v0.7.2

paper-ktft-v1.1

paper-ktft-v1

v0.7.1

v0.7.0.post4

v0.7.0.post1

v0.7.0

v0.6.4

v0.6.3.post1

v0.6.3

v0.6.2.post3

v0.6.2.post2

v0.6.2.post1

v0.6.2

v0.6.1

v0.6.1.post1

v0.5.3

v0.5.2.post1

v0.5.1

v0.5.0.post1

v0.5.0

v0.4.4

v0.4.3

v0.4.2

v0.4.1

v0.3.2

v0.3.1

v0.3

v0.2.4post1

v0.2.4

v0.2.3post2

v0.2.3post1

v0.2.3

v0.2.2

v0.2.2rc2

v0.2.2rc1

v0.2.1.post1

v0.2.1

v0.2.0

v0.1.4

v0.1.3

v0.1.2

v0.1.1

v0.1.0

Labels

Clear labels   

bug

  

duplicate

  

enhancement

  

good first issue

  

help wanted

  

pull-request

Mirrored from GitHub Pull Request

  

stale

No labels

bug

duplicate

enhancement

good first issue

help wanted

pull-request

stale

Milestone

Clear milestone

No items

No milestone

Projects

Clear projects

No items

No project

Assignees

Clear assignees

No assignees

1 participant

Notifications

Subscribe

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

kvcache-ai/ktransformers#38

No description provided.