[Bugfix] Fix DP/EP Shared Expert With Monolithic Kernels

[robertgshaw2-redhat]

Summary

• previously, we gave the shared expert to the MK if A2A is used, we should only do this is the MK can support Shared Expert overlap. I updated this to happen if we have DeepEP LL which supports this
• fixes: [Bug]: DeepSeek v3.2 FP8 Failure to start server #36050

[gemini-code-assist]

Code Review

This pull request aims to fix an issue with shared expert computation in monolithic kernels by restricting the passing of shared_experts to the FusedMoEKernel only when the deepep_low_latency All-to-All backend is used, as it's the only one that supports shared expert overlap.

While the intent is correct, the change appears to introduce a critical bug. For other All-to-All backends, this change will cause the shared expert computation to be skipped entirely. This is because the DefaultMoERunner delegates shared expert computation to the modular kernel for all All-to-All configurations, but with this change, the kernel will no longer receive the shared expert module for non-deepep_ll backends. I've added comments with details on the issue and suggestions for a fix.

[gemini-code-assist]

This change correctly identifies that only the deepep_low_latency backend supports shared expert overlap within the modular kernel. However, this introduces a critical issue for other All-to-All backends (e.g., deepep_high_throughput, mori).

Here's the breakdown of the issue:

1. For any All-to-All backend, a FusedMoEKernel is created, so quant_method.mk_owns_shared_expert becomes True.
2. This prevents DefaultMoERunner from computing the shared experts, as it delegates this task to the modular kernel.
3. With this change, for non-deepep_ll backends, shared_experts is passed as None to the FusedMoEKernel.
4. Consequently, the FusedMoEKernel also doesn't compute the shared experts.

This results in the shared expert computation being skipped entirely for these configurations, likely leading to incorrect model outputs.

To fix this, the logic that determines whether the modular kernel "owns" the shared expert computation needs to be updated. For instance, DefaultMoERunner should handle shared experts if use_all2all_kernels is true but use_deepep_ll_kernels is false.

[gemini-code-assist]

Similar to the change in fp8.py, this modification correctly restricts passing shared_experts to the FusedMoEKernel to only when the deepep_low_latency backend is used. However, it creates the same critical issue for other All-to-All backends.

The shared expert computation will be skipped for backends like deepep_high_throughput because:

1. quant_method.mk_owns_shared_expert will be True, so DefaultMoERunner won't run the shared experts.
2. The FusedMoEKernel will receive shared_experts=None and will also not run them.

This logic needs to be reconciled to ensure shared experts are always computed. The DefaultMoERunner should likely handle the shared expert computation when a modular kernel is used but does not support shared expert overlap (i.e., when use_all2all_kernels is true but use_deepep_ll_kernels is false).

[mergify]

Hi @robertgshaw2-redhat, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[bnellnm]

Should this condition be prepare_finalize.supports_async? That's the only time it really matters for the MK to call shared_experts.

[ProExpertProg]

Good with this if @bnellnm is