feat(kt-kernel): add score-aware expert placement

[githubshaurya]

Summary

This PR adds an optional score_aware_layer_balanced strategy for generating GPU expert masks in kt-kernel.

The existing default behavior is unchanged:

generate_gpu_experts_masks(activation_freq, num_gpu_experts)

still uses the current global frequency top-k placement.

The new strategy adds:

• EMA-smoothed activation scores via previous_scores and alpha
• optional min_experts_per_layer / max_experts_per_layer
• deterministic tie-breaking
• optional placement report with selected-score hit rate and per-layer expert counts

Motivation

KTransformers already supports [CPU-GPU expert scheduling] and documents placement strategies such as uniform, frequency, front-loading, and random.

The current [[generate_gpu_experts_masks](https://github.com/kvcache-ai/ktransformers/blob/main/kt-kernel/python/experts_base.py)] helper uses global top-k frequency placement. That is useful, but it can over-concentrate GPU experts in a few layers. This PR keeps the existing behavior intact while adding a score-aware, layer-balanced option for offline profiling and future scheduling work.

This also aligns with the [KTransformers 2026 Q2 roadmap](#1921), which mentions expert offloading/scheduling optimization and CPU-GPU coordination.

Background

This PR is inspired by recent hybrid MoE inference work:

• [HybriMoE]: motivates score-aware expert placement from changing expert activation patterns.
• [PreScope]: motivates layer-aware scheduling under resource-constrained MoE inference.

This PR does not implement a full runtime scheduler, prefetcher, cache manager, CUDA kernel, or SGLang runtime change. It only adds a small placement-mask utility while preserving default behavior.

Changed files

• kt-kernel/python/expert_placement.py

  • Adds a pure-Python expert placement planner.

• kt-kernel/python/experts_base.py

  • Extends generate_gpu_experts_masks while preserving default behavior.

• kt-kernel/python/__init__.py

  • Exports plan_gpu_expert_placement.

• kt-kernel/test/per_commit/test_expert_placement.py

  • Adds CPU-only tests for compatibility, EMA scoring, layer bounds, determinism, reports, and invalid inputs.

• doc/en/kt-kernel/experts-sched-Tutorial.md

  • Documents the new low-level utility.

Validation

python -m py_compile \
  kt-kernel/python/expert_placement.py \
  kt-kernel/python/experts_base.py \
  kt-kernel/python/__init__.py \
  kt-kernel/test/per_commit/test_expert_placement.py

PYTHONPATH=kt-kernel/python pytest -q kt-kernel/test/per_commit/test_expert_placement.py

PYTHONPATH=kt-kernel/python pytest -q kt-kernel/test/test_generate_gpu_experts_masks.py

Scope

This PR is intentionally limited to placement-mask generation. It avoids CUDA, kernel changes, SGLang runtime changes, and runtime prefetch/cache behavior.

[gemini-code-assist]

Code Review

This pull request introduces a new expert placement planning module (expert_placement.py) for heterogeneous MoE inference, adding a score_aware_layer_balanced strategy with EMA-smoothed scores and per-layer bounds, and integrates it into the existing generate_gpu_experts_masks API. Feedback focuses on PyTorch performance optimizations, specifically converting tensors to lists to avoid slow .item() calls inside nested loops, and vectorizing the minimum layer coverage constraint loop using torch.topk and scatter_.

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]

Calling .item() on individual tensor elements inside a nested loop is extremely slow in PyTorch due to host-device synchronization and Python-C++ boundary crossing overhead. Converting the tensors to Python lists using .tolist() beforehand allows the loop to run entirely in pure Python, which is orders of magnitude faster.

Suggested change

	candidates: list[tuple[float, int, int]] = []
	for layer_idx in range(num_layers):
	for expert_idx in range(num_experts_per_layer):
	if not bool(mask[layer_idx, expert_idx].item()):
	candidates.append(
	(float(scores[layer_idx, expert_idx].item()), layer_idx, expert_idx)
	)
	scores_list = scores.tolist()
	mask_list = mask.tolist()
	candidates: list[tuple[float, int, int]] = []
	for layer_idx in range(num_layers):
	for expert_idx in range(num_experts_per_layer):
	if not mask_list[layer_idx][expert_idx]:
	candidates.append(
	(scores_list[layer_idx][expert_idx], layer_idx, expert_idx)
	)

[gemini-code-assist]

The minimum layer coverage constraint can be fully vectorized using torch.topk along dim=1 and Tensor.scatter_, completely avoiding the Python loop over num_layers.

    # First satisfy the minimum layer coverage constraint.
    if min_experts_per_layer > 0:
        _, expert_indices = torch.topk(
            scores,
            k=min_experts_per_layer,
            dim=1,
            largest=True,
            sorted=False,
        )
        mask.scatter_(1, expert_indices, True)
        per_layer_count = [min_experts_per_layer] * num_layers