Fix dpsk-r1-fp4 tp8 by reverting two commits (#13162 and #13341)

[Qiaolin-Yu]

Motivation

Before this pr, this command will hang/fail.

SGLANG_ENABLE_SPEC_V2=1 /root/.python/sglang/bin/python -m sglang.launch_server --model-path nvidia/DeepSeek-R1-0528-FP4-v2 --trust-remote-code --quantization modelopt_fp4 --tp 8  --speculative-algorithm=EAGLE  --port 40020   --kv-cache-dtype fp8_e4m3   --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 8}'

Modifications

Accuracy Tests

Benchmarking and Profiling

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.

[gemini-code-assist]

Summary of Changes

Hello @Qiaolin-Yu, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request primarily focuses on resolving a bug related to dpsk-r1-fp4 tp8 by reverting specific prior commits. Concurrently, it refines the expert location and mapping mechanisms within the system and introduces a new tile_tokens_dim parameter across several MoE and quantization layers, likely to align with the corrected behavior or future optimizations.

Highlights

• Bug Fix via Reversion: This pull request addresses an issue with dpsk-r1-fp4 tp8 by reverting two specific commits (Fix nan in global scaling factor for large scale nvfp4 EP #13162 and Clean up deprecated tile_tokens_dim for next flashinfer #13341), aiming to restore stability or correct unintended behavior introduced by those changes.
• Expert Location Logic Simplification: The logical_to_all_physical method in expert_location.py has been simplified by removing the require_global_experts parameter and its associated conditional logic, streamlining how physical experts are identified.
• Refined Physical Expert Mapping: The logic for mapping logical experts to physical experts has been improved to prevent redundant assignments, ensuring that a nearest_expert is only added if it's not already present in the mapped list.
• Introduction of tile_tokens_dim Parameter: A new parameter, tile_tokens_dim=None, has been introduced and passed to various MoE and quantization layer methods (fused_moe_triton/layer.py, fp8.py, modelopt_quant.py, mxfp4.py), indicating a potential new configuration option or a cleanup related to the reverted changes.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request appears to be a revert to fix issues with DeepSeek models, specifically related to expert location mapping and MoE kernel parameters.

The changes to expert_location.py aim to correct a bug in how the nearest physical expert is handled. However, the new logic introduces a different issue where the condition to update the expert list is never met, making the optimization ineffective. I've provided a suggestion to fix this.

The other changes consistently add the tile_tokens_dim parameter to MoE kernel calls across several files. This seems to be a necessary fix to align with an underlying API change, likely from the reverted commits. These changes look correct.

Overall, the PR is a step in the right direction, but the expert location logic needs a further fix to work as intended.

[gemini-code-assist]

The condition nearest_expert not in mapped_physical_experts will always be False when nearest_expert != -1. This is because _find_nearest_expert is called with mapped_physical_experts as candidate_physical_expert_ids, and it returns an element from that list (or -1). As a result, this if block is effectively dead code, and the nearest expert is never prioritized.

This defeats the purpose of finding the nearest expert for locality optimization.

To correctly prioritize the nearest expert, it should be moved to the front of the mapped_physical_experts list. A swap with the first element is a safe way to do this while preserving the list length and avoiding duplicates.

Suggested change

	mapped_physical_experts = logical_to_all_physical_map[layer_id][
	logical_expert_id
	]
	if (
	nearest_expert != -1
	and nearest_expert not in mapped_physical_experts
	):
	mapped_physical_experts[0] = nearest_expert
	mapped_physical_experts = logical_to_all_physical_map[layer_id][
	logical_expert_id
	]
	if nearest_expert != -1:
	# The `_find_nearest_expert` function returns an expert from `mapped_physical_experts`.
	# We should move it to the front to prioritize it for dispatch.
	try:
	current_idx = mapped_physical_experts.index(nearest_expert)
	if current_idx > 0:
	# Swap with the first element to preserve list length and avoid duplicates.
	mapped_physical_experts[0], mapped_physical_experts[current_idx] = (
	mapped_physical_experts[current_idx],
	mapped_physical_experts[0],
	)
	except ValueError:
	# This path should ideally not be taken.
	# It implies `_find_nearest_expert` returned an expert not in the candidate list.
	# As a fallback, we replace the first expert.
	mapped_physical_experts[0] = nearest_expert

[Kangyan-Zhou]

Tested B200 CI on a different machine works, will merge this PR for now