[feat](kt-kernel): add co-activation-aware and online-EMA GPU expert placement

[neerajdad123-byte]

What does this PR do?

The default expert placement (generate_gpu_experts_masks) keeps the globally
most-frequent experts, scoring each independently. That is blind to which experts
co-fire on the same token: if two experts almost always route together but only one
is GPU-resident, every such token still pays the slow CPU path for the other.

This adds kt-kernel/python/expert_placement.py with two drop-in capabilities that
raise the per-token GPU hit rate at the same GPU-expert budget:

• plan_gpu_expert_placement — static co-activation-aware planner with coverage
  (whole-token set-cover) and cluster (greedy co-firing cluster) strategies.
  Degrades to frequency top-k when only marginal counts are available, so it never
  regresses the existing behaviour.
• EmaHotnessTracker — online per-layer EMA of expert activation that emits an
  up-to-date residency mask as the live workload shifts, in O(top_k) work per layer
  per token. Model- and workload-agnostic (no offline profile). Optionally warm-starts
  from activation_freq so the cold-start mask equals frequency top-k and never regresses.

Both produce the same (num_layers, num_experts) boolean mask the pipeline already
consumes, so integration is drop-in. Pure Python/torch — no routing, no weight movement.
Also adds gpu_hit_rate and token_resident_rate metrics.

Fixes # (N/A — enhancement)

Before submitting

• Did you read the contributor guideline?
• Did you write any new necessary tests? (30 CPU-only tests in
  test/per_commit/test_expert_placement.py, all passing; black clean)

[gemini-code-assist]

Code Review

This pull request introduces a co-activation-aware GPU expert placement planner (expert_placement.py) for Mixture of Experts (MoE) inference, along with comprehensive unit tests. The planner aims to maximize whole-token GPU residency using greedy clustering and weighted maximum-coverage strategies, and includes an online EMA tracker (EmaHotnessTracker) to adapt to shifting workloads. The review feedback identifies several critical performance bottlenecks where Python loops are used over token traces (in build_coactivation_matrix, observe_batch, token_resident_rate, and _coverage_select) and provides vectorized PyTorch suggestions to optimize execution speed.

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 current implementation of build_coactivation_matrix uses a Python loop over all tokens (for token_idx in range(num_tokens)). For large traces, this Python loop will introduce a significant performance bottleneck. We can fully vectorize this operation using PyTorch's matrix multiplication (torch.mm) and scatter_ to construct a binary indicator matrix, which runs entirely in optimized C++.

    num_tokens, top_k = traces.shape
    clamped_traces = torch.where((traces >= 0) & (traces < num_experts), traces, num_experts)
    indicator = torch.zeros(num_tokens, num_experts + 1, dtype=torch.float32, device="cpu")
    indicator.scatter_(1, clamped_traces, 1.0)
    indicator = indicator[:, :num_experts]
    coactivation = torch.mm(indicator.t(), indicator)
    return coactivation

[gemini-code-assist]

The observe_batch method currently processes tokens sequentially in a Python loop (for token_ids in ids:). For online serving with larger batch sizes, this loop will be extremely slow. We can vectorize the batch EMA update mathematically by calculating the exponential decay weights for each token in the batch and performing a single vectorized index_add_ operation.

Suggested change

	for layer_idx, traces in token_traces.items():
	ids = traces.to(device="cpu", dtype=torch.long)
	if ids.dim() == 1:
	ids = ids.unsqueeze(0)
	for token_ids in ids:
	self.observe(layer_idx, token_ids)
	for layer_idx, traces in token_traces.items():
	if layer_idx < 0 or layer_idx >= self.num_layers:
	continue
	ids = traces.to(device="cpu", dtype=torch.long)
	if ids.dim() == 1:
	ids = ids.unsqueeze(0)
	num_tokens, top_k = ids.shape
	if num_tokens == 0:
	continue
	row = self._hotness[layer_idx]
	row.mul_(self.decay ** num_tokens)
	exponent = torch.arange(num_tokens, dtype=torch.float32, device="cpu")
	weights = (1.0 - self.decay) * torch.pow(self.decay, num_tokens - 1 - exponent)
	flat_weights = weights.unsqueeze(1).expand(-1, top_k).flatten()
	flat_ids = ids.flatten()
	valid_mask = (flat_ids >= 0) & (flat_ids < self.num_experts)
	if valid_mask.any():
	row.index_add_(0, flat_ids[valid_mask], flat_weights[valid_mask])

[gemini-code-assist]

The token_resident_rate function uses a Python loop over all tokens (for token_ids in ids:) to check if all routed experts are resident. This can be fully vectorized using PyTorch's element-wise logical operations, avoiding the slow Python loop entirely.

Suggested change

	for token_ids in ids:
	valid = token_ids[(token_ids >= 0) & (token_ids < num_experts)]
	if valid.numel() == 0:
	continue
	total_tokens += 1
	if bool(layer_mask[valid].all()):
	fully_resident += 1
	valid = (ids >= 0) & (ids < num_experts)
	has_valid = valid.any(dim=1)
	total_tokens += int(has_valid.sum().item())
	clamped_ids = ids.clamp(min=0, max=num_experts - 1)
	resident = layer_mask[clamped_ids]
	is_ok = resident | ~valid
	token_ok = is_ok.all(dim=1)
	fully_resident += int((token_ok & has_valid).sum().item())

[gemini-code-assist]

In _coverage_select, token_sets is iterated over in every step of the while loop to compute the missing experts. As more experts are selected, many tokens become fully covered. We can optimize this by dynamically filtering out fully covered tokens from token_sets at the end of each iteration, which significantly reduces the number of tokens processed in subsequent iterations.

Suggested change

	for e in best_key:
	selected[e] = True
	selected_set.add(e)
	for e in best_key:
	selected[e] = True
	selected_set.add(e)
	token_sets = [s for s in token_sets if not s.issubset(selected_set)]