[perf](kt-kernel): SFT end2end

[yyj6666667]

Depends on #2086.

This is a stacked PR based exactly on the #2086 head commit 66f5f15. The implementation specific review commit is 6eeffa7.

What does this PR do?

Makes the BF16 buffers managed by KT C++ the single authoritative optimizer-gradient storage for CPU-only AMXBF16_SFT Full-FT, LoRA, and Hybrid training.

• Binds Full-FT and fused/PEFT LoRA Parameter.grad directly to the preallocated KT buffers.
• Returns no base-weight gradient from KTMoEFunction.backward, removing the duplicate PyTorch AccumulateGrad and large aten::add_ path.
• Uses overwrite on the first microbatch in an optimizer window and in-place C++ accumulation on later microbatches.
• Releases authoritative .grad aliases after optimizer.step; Full-FT gradient buffers use torch.empty, and C++ clears them before first use.
• Adds optimizer-gradient scaling, pointer/failure recovery, TP reduce accumulation, and overwrite/accumulate/lazy-clear profiler counters.
• Keeps rank 0 authoritative in distributed execution and applies 1 / world_size only to parameter gradients.
• Unifies Full, LoRA, and Hybrid lifecycle validation, including synchronous and asynchronous failure handling.

The scope remains CPU-only AMXBF16_SFT. Quantized backends, GPU experts, and SkipLoRA behavior are unchanged.

Validation

• Python authoritative-gradient lifecycle: 12/12 tests passed, plus non-reentrant checkpoint reuse coverage.
• Full, LoRA, and Hybrid with GAS 1/2/4: all passed against FP32 references.
• TP1/TP2 consistency: all matrix cases passed.
• Two-rank rank-0 ownership tests: TP1 and TP2 passed without collective hangs.
• BF16 dWeight overwrite, accumulation, scaling, tail, and stride tests passed on AVX512-BF16 hardware.
• C++ extension and test_bf16_dweight targets build successfully.
• Qwen3 MoE 15-step Full-FT smoke tests had finite loss/gradients and changed KT parameters.
• Torch profiler trace contains no KT base-weight AccumulateGrad, large aten::add_, or KT zero-gradient scan.

Before submitting

• Read the contributor guideline.
• Added necessary Python and C++ tests.

[gemini-code-assist]

Code Review

This pull request introduces full fine-tuning (full weight gradient) support and lightweight staged profiling for the KT SFT MoE operators. Key changes include the implementation of BF16DWeightKernel for outer-product gradient accumulation, direct BF16 weight reloading, and authoritative optimizer-gradient lifecycle management to reduce memory overhead and avoid full object recreation. It also adds a staged profiler (SFTProfiler) and comprehensive unit tests. The review feedback highlights a bug in multi-threaded profiling where thread_local timings are incorrectly read on the main thread, a conceptual indexing issue in from_bb_transposed regarding tile strides, and an optimization opportunity in sync_kt_lora_gradients to avoid an unnecessary temporary CPU tensor allocation during cross-device copies.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

The amx::bf16_dweight_timings() function returns a thread_local reference. Since the finalizer lambda [this](int _) is executed on the main thread after all tasks are completed, timings will only contain the metrics for the tasks executed by the main thread itself, completely ignoring the metrics from other worker threads in the pool. This leads to severely underreported profiling statistics in multi-threaded environments. Since SFTProfiler::record_ns is already thread-safe (using std::atomic::fetch_add), a much simpler and correct approach is to record the timings directly from the worker tasks (e.g., at the end of each task or periodically), or accumulate them into a shared/atomic structure.

[gemini-code-assist]

In from_bb_transposed, the destination tile dst_tile has dimensions K_STEP (rows) by N_STEP (columns). Therefore, its row stride is N_STEP. The indexing for dst_tile should use j * N_STEP + i instead of j * K_STEP + i. Although K_STEP and N_STEP are currently both 32, using N_STEP is conceptually correct and prevents potential out-of-bounds or layout corruption bugs if the step sizes ever differ in future configurations.

          for (int i = 0; i < N_STEP; i++) {
            for (int j = 0; j < K_STEP; j++) {
              dst_tile[j * N_STEP + i] = src_tile[i * K_STEP + j];
            }
          }

[gemini-code-assist]

In sync_kt_lora_gradients, when copying the synchronized base weight gradients back to the CPU buffer, grad_buf.copy_(grad_gpu.cpu()) is used. This creates an unnecessary temporary CPU tensor via .cpu(). Since PyTorch's copy_ operator natively supports cross-device copies, you can directly use grad_buf.copy_(grad_gpu) to avoid the overhead of allocating and copying to a temporary CPU tensor, matching the optimized pattern used for LoRA gradients on line 783.

Suggested change

	if grad_buf is not None:
	grad_gpu = grad_buf.cuda()
	dist.all_reduce(grad_gpu, op=dist.ReduceOp.SUM)
	grad_gpu.div_(world_size)
	grad_buf.copy_(grad_gpu.cpu())
	if grad_buf is not None:
	grad_gpu = grad_buf.cuda()
	dist.all_reduce(grad_gpu, op=dist.ReduceOp.SUM)
	grad_gpu.div_(world_size)
	grad_buf.copy_(grad_gpu)