[aoshen02]

Purpose

#64583 turned off the Kimi-K3 latent-MoE tail fusion whenever sleep mode is on, and logged a warning. At the time, level-2 sleep discarded the whole weights pool. The fusion's collective kernel keeps its state in that pool: the Lamport flags (_routed_flags, _shared_flags) and the peer pointer table. None of these are parameters, so they were lost on wake.

Since #66010, level 2 discards only allocations that hold nothing but parameters, and offloads and restores everything else at the same address. That state now survives sleep, so this PR removes the fallback and its warning. The tail fusion now runs with sleep mode on, at TP 8 and 16 on SM100.

• latent_moe_runner.py: delete the sleep-mode fallback (11 lines).
• test_latent_moe_tail.py: delete test_sleep_mode_falls_back_to_default_path and the enable_sleep_mode parameter that only it used.

Test Plan

Kimi-K3 at TP 8 across two GB200 nodes (multi-node NVLink), with dummy weights and 4 layers (hf_overrides), CUDA graphs on, enable_sleep_mode=True and sleep_preserve_parameter_names=["*"]. Keeping every parameter means that any difference after wake comes from non-parameter state.

The run checks each rank's LatentMoERunners through collective_rpc, then generates greedily for 16 prompts and 48 tokens:

• before sleep, twice;
• after sleep(1) + wake_up();
• after each of two sleep(2) + wake_up(["weights"]) + wake_up(["kv_cache"]) cycles.

Each output is compared with the first, token for token.

Negative control: the same run, but after the first level-2 wake the collective kernel's Lamport flags are zeroed. This emulates what the old level 2 did to them.

tests/models/kimi_k3/test_latent_moe_tail.py also runs on GB200.

Test Result

Run	Tail fusion on (runners per rank)	Repeat	After L1	After L2 #1	After L2 #2
This PR	3 / 3 on all 8 ranks, fused path taken	16/16	16/16	16/16	16/16
Negative control (flags zeroed after L2 #1)	3 / 3 on all 8 ranks			hangs: RPC call to sample_tokens timed out

The control shows that the test catches lost tail-fusion state. Without #66010 (main at 2ced9ce with the same fallback removed), the same run crashes at the first level-2 wake with cudaErrorIllegalAddress, even though the fusion was not confirmed to be on there.

test_latent_moe_tail.py: 1 passed, 3 skipped (the skipped kernel tests need more GPUs).

Not covered: real Kimi-K3 weights and GSM8K (no checkpoint available here), TP 16, and level 2 with a real weight update. In that path, parameters are discarded and rewritten. That is the generic #66010 path, which this PR doesn't change.

---

Pull Request Checklist

• AI assistance was used during the creation of this PR.
• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #100772 for commit a95d224dfeb7.