[Disagg] Fix MegaMoE topk_ids dtype mismatch and FakeKVManager missing kv_args

[JoeLee314]

Motivation

When serving DeepSeek-V4-Pro with MegaMoE enabled (SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE=1) and disaggregation decode mode (--disaggregation-transfer-backend fake --ep-dispatch-algorithm fake), the server crashes during initialization with two independent errors:

Bug 1 — dtype mismatch in mega_moe_pre_dispatch:

tvm.error.InternalError: Tensor match failed for Tensor<32, 6>[strides=<6, 1>, dtype=int64, device=cuda:0]
  at mega_moe_pre_dispatch.cuh:150
- Root cause: Dtype value [int64] not in the allowed options: [int32]

The topk_output.topk_ids is int64 with -ep-dispatch-algorithm fake, but the JIT-compiled mega_moe_pre_dispatch kernel only accepts int32.
Here is the analys:

1. When --ep-dispatch-algorithm is not set, ExpertLocationDispatchInfo.init_new() returns None (expert_location_dispatch.py:50), so topk_ids_logical_to_physical() returns topk_ids as-is — int32 is preserved.
2. When --ep-dispatch-algorithm fake is set, it routes to _topk_ids_logical_to_physical_dynamic() (expert_location_dispatch.py:93), which does:
   topk_ids = info.partial_logical_to_all_physical_map[topk_ids, chosen_dispatch_index]
3. partial_logical_to_all_physical_map is created via torch.tensor() without an explicit dtype, defaulting to int64 (expert_location_dispatch.py:422). In PyTorch, tensor indexing result dtype follows the source tensor, not the index tensor — so int64_tensor[int32_idx] returns int64, silently undoing the earlier .to(torch.int32) cast.
4. The int64 topk_ids then flows into mega_moe_pre_dispatch which only accepts int32, triggering the TVM dtype check error.

Bug 2 — FakeKVManager missing kv_args attribute:

AttributeError: 'FakeKVManager' object has no attribute 'kv_args'

FakeKVManager inherits from BaseKVManager but does not store the args parameter as self.kv_args, unlike CommonKVManager. The disaggregation decode event loop accesses self.kv_manager.kv_args.state_types, which crashes at runtime.

Both bugs are hit when running:

SGLANG_OPT_USE_DEEPGEMM_MEGA_MOE=1 \
SGLANG_OPT_FIX_HASH_MEGA_MOE=1 \
SGLANG_OPT_FIX_MEGA_MOE_MEMORY=1 \
SGLANG_OPT_FIX_NEXTN_MEGA_MOE=1 \
SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK=8320 \
SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256 \
sglang serve \
    --trust-remote-code \
    --model-path DeepSeek-V4-Pro \
    --tp 8 --dp 8 --enable-dp-attention \
    --moe-a2a-backend deepep \
    --disaggregation-mode decode \
    --ep-dispatch-algorithm fake \
    --disaggregation-transfer-backend fake \
    --dist-init-addr 127.0.0.1:30435 \
    --mem-fraction-static 0.85 \
    --host 0.0.0.0 --port 30001

Modifications

python/sglang/srt/layers/moe/mega_moe.py — _run_mega_routed:

• Cast topk_ids to int32 before passing to mega_moe_pre_dispatch, since the kernel only accepts int32 tensors. The else branch already creates an int32 empty tensor, so this makes the two branches consistent.
  The alternative would be fixing _topk_ids_logical_to_physical_dynamic() or the map dtype in expert_location_dispatch.py. However, casting at the MegaMoE entry point (mega_moe.py) is a more robust choice because:
  • It makes MegaMoE self-contained — the kernel's contract (int32 input) is enforced at its own boundary regardless of upstream dtype changes.
  • Multiple topk paths (grouped, biased, fused, custom) converge here; fixing any single upstream path would miss others.
  • It avoids risking regressions in non-MegaMoE paths where int64 topk_ids may be used by other kernels without issue.

python/sglang/srt/disaggregation/fake/conn.py — FakeKVManager.__init__:

• Add self.kv_args = args, matching the pattern already used in CommonKVManager. This is needed because the decode event loop accesses self.kv_manager.kv_args.state_types during preallocation.

Accuracy Tests

Not applicable — the first fix is a dtype cast that preserves values (int64→int32 is safe since expert indices are small), and the second fix stores a reference that was already being accessed at runtime.

Speed Tests and Profiling

Not applicable — both fixes are one-line changes with no impact on the forward compute path. The .to(torch.int32) cast is a no-op when the tensor is already on GPU and the values fit in int32.

Verified that the server starts successfully and prints The server is fired up and ready to roll! with the above command on 8× NVIDIA B200 GPUs.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Missing run-ci label — add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked — run-ci is required first.

[gemini-code-assist]

Code Review

This pull request adds a state assignment in the disaggregation connection and casts top-k IDs to int32 in the Mega MoE layer for better type consistency. Feedback suggests also casting top-k weights to float32 to ensure compatibility with JIT-compiled kernels and consistency with the empty-token branch.

[ShangmingCai]

disaggregation fake backend modification LGTM

[JoeLee314]

Thanks for the review! I just squashed the commits into one and force-pushed, so the previous approval got cleared. Could you re-approve when you get a chance? No code changes, just a clean single commit now.

[ShangmingCai]

Actually, I haven't approved yet since mega_moe part still requires a review from other reviewers. I am not an expert on mega moe, so you should ping other reviewers to check this fix, then we can run the CI.

[JoeLee314]

Thanks for the clarification! Got it.

[Fridge003]

/rerun-test test_deepseek_v4_flash_fp4_megamoe_b200.py

[github-actions]

⛔ test_deepseek_v4_flash_fp4_megamoe_b200.py: Dispatch failed: 422

[Fridge003]

/rerun-test test_deepseek_v4_flash_fp4_megamoe_b200.py

[github-actions]

⛔ test_deepseek_v4_flash_fp4_megamoe_b200.py: Dispatch failed: 422

[Fridge003]

/rerun-test test_deepseek_v4_flash_fp4_megamoe_b200.py

[github-actions]

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/dsv4/test_deepseek_v4_flash_fp4_megamoe_b200.py