Skip to content

[feat](kt-kernel): add native MXFP4 routed-expert LoRA SFT - #2185

Open
yyj6666667 wants to merge 2 commits into
kvcache-ai:mainfrom
yyj6666667:yyj/dsv4-mxfp4-routed-sft-kernel
Open

[feat](kt-kernel): add native MXFP4 routed-expert LoRA SFT#2185
yyj6666667 wants to merge 2 commits into
kvcache-ai:mainfrom
yyj6666667:yyj/dsv4-mxfp4-routed-sft-kernel

Conversation

@yyj6666667

@yyj6666667 yyj6666667 commented Sep 3, 2026

Copy link
Copy Markdown
Collaborator

What does this PR do?

This PR adds standalone CPU LoRA forward/backward support for DeepSeek V4 native MXFP4 routed experts.

Main changes:

  • Consume the checkpoint-native E2M1 weights and group-32 UE8M0 scales without creating a persistent dense BF16/FP32 base copy.
  • Compute base dX directly from packed row-major weights, without an additional transposed base-weight copy.
  • Implement the DeepSeek V4 asymmetric SwiGLU clamp and its backward derivative, with LoRA applied before the clamp.
  • Reuse KT authoritative LoRA gradients: overwrite on the first microbatch, accumulate afterward, and lazy-clear between optimizer windows.
  • Add source-level, compiled numerical, lifecycle, and real-checkpoint layer validation.

The scope of this PR is one CPU routed-expert MoE layer. Whole-model DeepSeek V4 attention backward and training-framework integration are left for a follow-up.

Related to #1921.

Tests

  • Python reference/API/gradient lifecycle: 49 passed.
  • Compiled TP1/TP2 numerical tests: GAS 1/2/8, different active-expert sets, lazy clear, checkpoint-cache replacement, and legacy MXFP4 inference regression passed.
  • Synthetic comparison against an FP32 reference: worst relative L2 0.008501, minimum cosine 0.999965.
  • Real DeepSeek-V4-Flash-0731 layer 3 smoke: 256 experts, H=4096, I=2048, top-k=6, rank=8; TP1 and TP2 passed.
  • Real-layer comparison: worst relative L2 0.007374, minimum cosine 0.999973; inactive expert gradients remained zero and the packed base checksum was unchanged after backward and a LoRA update.
  • CPU-only AVX512-BF16 extension rebuilt successfully on dual AMD EPYC 9355.

Before submitting

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant