[Feature]: Expert Parallelism with CompressedTensor FP4 or NVFP4

[Victor49152]

🚀 The feature, motivation and pitch

In order to further improve the performance of large models like Qwen3-VL-235B, fp4 plus expert parallel is a necessary feature.

Alternatives

No response

Additional context

The current fp4 checkpoint is quantized through llm-compressor.

The current quantization method class has assertion that blocks using --enable-expert-parallel.
AssertionError: Expert Parallelism / expert_map is currently not supported for CompressedTensorsW4A4MoeMethod.

Similar warning has been seen on NVFP4 method through modelopt.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[wangshangsam]

I was just reading the code and a couple of things seemed worth investigation to me.

1. Would FlashInfer (with fused_experts or flashinfer_cutlass_moe_fp4) support EP?

This error check happens in CompressedTensorsW4A4MoeMethod.apply which seems to be where kernel dispatching happens during a forward pass:

vllm/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

Lines 464 to 554 in aef368a

	if self.use_marlin:
	assert self.fused_experts is None
	return fused_marlin_moe(
	x,
	layer.w13_weight,
	layer.w2_weight,
	None,
	None,
	layer.w13_weight_scale,
	layer.w2_weight_scale,
	router_logits,
	topk_weights,
	topk_ids,
	global_scale1=layer.w13_weight_scale_2,
	global_scale2=layer.w2_weight_scale_2,
	quant_type_id=scalar_types.float4_e2m1f.id,
	apply_router_weight_on_input=apply_router_weight_on_input,
	global_num_experts=global_num_experts,
	expert_map=expert_map,
	workspace=layer.workspace,
	)
	elif self.fused_experts is not None:
	assert is_valid_flashinfer_cutlass_fused_moe(
	x, layer.w13_weight, layer.w2_weight
	), "Flashinfer CUTLASS Fused MoE not applicable!"
	return self.fused_experts(
	hidden_states=x,
	w1=layer.w13_weight,
	w2=layer.w2_weight,
	topk_weights=topk_weights,
	topk_ids=topk_ids,
	inplace=False, # TODO(shuw): fix later, now output is high prec
	activation=activation,
	global_num_experts=global_num_experts,
	expert_map=expert_map,
	apply_router_weight_on_input=apply_router_weight_on_input,
	)
	# FlashInfer fused experts path
	elif self.allow_flashinfer:
	from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
	flashinfer_cutlass_moe_fp4,
	)
	assert is_valid_flashinfer_cutlass_fused_moe(
	x, layer.w13_weight, layer.w2_weight
	), "Flashinfer CUTLASS Fused MoE not applicable!"
	assert self.moe_quant_config is not None
	return flashinfer_cutlass_moe_fp4(
	hidden_states=x,
	w1=layer.w13_weight,
	w2=layer.w2_weight,
	topk_weights=topk_weights,
	topk_ids=topk_ids,
	quant_config=self.moe_quant_config,
	inplace=False, # TODO(shuw): fix later, now output is high prec
	activation=activation,
	global_num_experts=global_num_experts,
	expert_map=expert_map,
	apply_router_weight_on_input=apply_router_weight_on_input,
	)
	else:
	from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4
	assert expert_map is None, (
	"Expert Parallelism / expert_map "
	"is currently not supported for "
	"CompressedTensorsW4A4MoeMethod."
	)
	assert self.moe_quant_config is not None
	# Cutlass moe takes in activations in BF16/Half precision
	# and fp4 quantized weights loaded from the checkpoint
	return cutlass_moe_fp4(
	a=x,
	w1_fp4=layer.w13_weight,
	w2_fp4=layer.w2_weight,
	topk_weights=topk_weights,
	topk_ids=topk_ids,
	quant_config=self.moe_quant_config,
	apply_router_weight_on_input=apply_router_weight_on_input,
	# TODO(bnell): derive these from arguments
	m=x.shape[0],
	n=layer.w2_weight.shape[2] * 2,
	k=x.shape[1],
	e=layer.w13_weight.shape[0],
	).to(x.dtype)

The overall logic of the above code block is:

 if self.use_marlin:
    # Marlin seems to be for backward compatibility purpose on 
    # older-generation-than-Blackwell GPUs that don't have nvfp4 tensor cores. 
    # This isn't relevant to our use case.
    ...
  elif self.fused_experts is not None:
    # FlashInfer's cutlass implementation wrapped in vLLM's `FusedMoEModularKernel` code path
    # for specific All2All backends (`export VLLM_ALL2ALL_BACKEND=deepep_high_throughput`
    # or `deepep_low_latency` or `pplx` or `flashinfer_all2allv`). This seems to be for EP + DP.
    ...
  elif elif self.allow_flashinfer:
    # Direct call into FlashInfer's cutlass implementation. This seems to be for EP + TP.
    ...
  else:
    # The fallback, vLLM's own CUTLASS FP4 kernel which indeed doesn't support EP.
    ...

It looks like the execution landed in the fallback code path. It would be worthwhile to investigate if the two FlashInfer (i.e., intended) code paths work as intended. You can probably (forcefully) turn FlashInfer on with

export VLLM_USE_FLASHINFER_MOE_FP4=1

export VLLM_FLASHINFER_MOE_BACKEND=throughput 
# This should be the default. There's a `latency` FlashInfer mode that I'll get to below.

# In the case of EP + DP, you can select the All2All backend with
export VLLM_ALL2ALL_BACKEND=allgather_reducescatter  # Default
# OR choose from:
# - naive
# - pplx
# - deepep_high_throughput
# - deepep_low_latency
# - flashinfer_all2allv

2. Why the execution landed in the fallback CUTLASS FP4 kernel path in the first place?

To my undersstanding, even if you don't set export VLLM_USE_FLASHINFER_MOE_FP4=1 and export VLLM_FLASHINFER_MOE_BACKEND=throughput, vLLM is supposed to pick out the most suitable path based on the GPU device capability and the parallelism setting (and whether FlashInfer is indeed available). So the fact that, running a OOTB commandline from the Qwen3 VL recipe with EP turned out would land in the fallback path, probably indicates that there's some default selection logic that might need fixing.

3. It seems that the more appropriate default path (flashinfer_cutlass_moe_fp4()) was taken for the TP=8 case.

In the nsys traces that you shared, I seem to be seeing FlashInfer kernels actually being used,

this cutlass::device_kernel<flashinfer::gemm::DeviceGemmFp4GemmSm100___nv_bfloat16_128_256_256_4_4_1_2SM::Sm10x11xOnly kernel seems to be coming from flashinfer_cutlass_moe_fp4() which seems to indicate that the default code path selection for TP-only seems more correct. Hence, where the code path selection went wrong might be EP-specific.

[wangshangsam]

Side note RE export VLLM_FLASHINFER_MOE_BACKEND=latency

This needs a direct call into FlashInfer's flashinfer.fused_moe.trtllm_fp4_block_scale_moe which seems to be available through the ModelOptNvFp4FusedMoE layer implementation

vllm/vllm/model_executor/layers/quantization/modelopt.py

Lines 1621 to 1685 in be44450

	if (
	self.allow_flashinfer
	and self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
	):
	import flashinfer
	from vllm.model_executor.models.llama4 import Llama4MoE
	assert self.fused_experts is None
	a1_gscale = layer.w13_input_scale_quant
	(hidden_states_fp4, hidden_states_scale_linear_fp4) = (
	flashinfer.fp4_quantize(
	x,
	a1_gscale,
	is_sf_swizzled_layout=False,
	)
	)
	use_llama4_routing = (
	custom_routing_function is Llama4MoE.custom_routing_function
	)
	routing_method_type = flashinfer.RoutingMethodType.DeepSeekV3
	if use_llama4_routing:
	routing_method_type = flashinfer.RoutingMethodType.Llama4
	routing_bias = e_score_correction_bias
	if routing_bias is not None:
	routing_bias = routing_bias.to(torch.bfloat16)
	out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
	routing_logits=router_logits
	if use_llama4_routing
	else router_logits.to(torch.float32),
	routing_bias=routing_bias,
	hidden_states=hidden_states_fp4,
	hidden_states_scale=hidden_states_scale_linear_fp4.view(
	torch.float8_e4m3fn
	).flatten(),
	gemm1_weights=layer.gemm1_weights_fp4_shuffled.data,
	gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.data.view(
	torch.float8_e4m3fn
	),
	gemm1_bias=None,
	gemm1_alpha=None,
	gemm1_beta=None,
	gemm1_clamp_limit=None,
	gemm2_weights=layer.gemm2_weights_fp4_shuffled.data,
	gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.data.view(
	torch.float8_e4m3fn
	),
	gemm2_bias=None,
	output1_scale_scalar=layer.g1_scale_c.data,
	output1_scale_gate_scalar=layer.g1_alphas.data,
	output2_scale_scalar=layer.g2_alphas.data,
	num_experts=global_num_experts,
	top_k=top_k,
	n_group=num_expert_group if num_expert_group is not None else 0,
	topk_group=topk_group if topk_group is not None else 0,
	intermediate_size=layer.intermediate_size_per_partition,
	local_expert_offset=layer.ep_rank * layer.local_num_experts,
	local_num_experts=layer.local_num_experts,
	routed_scaling_factor=None,
	tile_tokens_dim=None,
	routing_method_type=routing_method_type,
	do_finalize=True,
	)[0]
	return out

but not the CompressedTensorsW4A4MoeMethod layer implementation. So either there's work on CompressedTensorsW4A4MoeMethod that needs to support this, or we need modelopt eventually after the Qwen3 VL MoE quantization issue and the Qwen3 MoE weight loading issue are fixed.

FYI @soodoshll @shengliangxu @Edwardf0t1 @dsikka

[wenscarl]

#27223 to address the code path issue. @wangshangsam The same applies to compressed tensors quantization. Just wip.

[wangshangsam]

Confirmed with @wenscarl that, as of today, those env vars are still needed to enable FlashInfer.

[Victor49152]

the more appropriate default path (flashinfer_cutlass_moe_fp4()) was taken for the TP=8

I took a closer look at the profile you shared. This flashinfer kernel seems to occur at attention part, which is expected since I was running with VLLM_ATTENTION_BACKEND=FLASHINFER.

Confirmed with @wenscarl that, as of today, those env vars are still needed to enable FlashInfer.

With env variables VLLM_USE_FLASHINFER_MOE_FP4=1 VLLM_FLASHINFER_MOE_BACKEND=throughput set, it goes to a different code path as you mentioned. And fused_experts goes to FlashinferExperts, in this case the EP works.