diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index b4aef340c..fb84b5c76 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -354,7 +354,7 @@ SGLang supports various environment variables that can be used to configure its SGLANG_ENABLE_MOE_DEFERRED_FINALIZE Defer the MoE finalize step to overlap it with other work. - false + true SGLANG_PATCH_TOKENIZER diff --git a/python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu b/python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu index 8ae8072b8..2a96c8f2a 100644 --- a/python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu +++ b/python/sglang/kernels/jit/csrc/moe/moe_finalize_fuse_shared.cu @@ -35,10 +35,10 @@ * ``DeepseekV3MoE.forward``, and gives the downstream allreduce+rmsnorm * a clean PDL handoff. * - * Expert-weight dtype is templated on ``TypeExpW`` so we support both the - * bf16 and fp32 topk-weight paths (DSv3/K2.5 trtllm backends use fp32 - * because their ``_routing_logits_dtype = torch.float32``; other backends - * use bf16). + * Expert-weight dtype is templated on ``TypeExpW`` so we accept both bf16 + * and fp32 topk weights. The trtllm deferred-finalize path always feeds bf16 + * (the trtllm-gen routing kernel emits bf16 for every routing method); fp32 + * is kept for callers that produce topk weights in fp32. * * Expert-weight scale convention: in our target backends * (flashinfer trtllm nvfp4 + unquantized), ``apply_routed_scaling_factor_on_output`` diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 99fa1f44b..eee06bdc2 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1332,7 +1332,7 @@ class Envs: # Sglang Cache Dir SGLANG_CACHE_DIR = EnvStr(os.path.expanduser("~/.cache/sglang")) SGLANG_FLASHINFER_AUTOTUNE_CACHE = EnvBool(True) - SGLANG_ENABLE_MOE_DEFERRED_FINALIZE = EnvBool(False) + SGLANG_ENABLE_MOE_DEFERRED_FINALIZE = EnvBool(True) # Plugin system SGLANG_PLATFORM = EnvStr("") diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index d02de06f2..f821d095c 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -1130,13 +1130,6 @@ def fused_experts_none_to_flashinfer_trtllm_fp4( result = trtllm_fp4_block_scale_moe(**moe_kwargs) if defer_finalize: gemm2_out, expert_weights, expanded_idx_to_permuted_idx = result[:3] - # FIXME(kpham-sgl): flashinfer sizes this buffer from routing_logits - # dtype (fp32 in DSv3 decode) but always writes bf16 weights into it. - # Reinterpret the live bf16 prefix. Fix upstream alloc to drop this, - # tracking in https://github.com/flashinfer-ai/flashinfer/issues/3595 - if expert_weights.dtype == torch.float32: - n, k = expert_weights.shape - expert_weights = expert_weights.view(torch.bfloat16).view(-1, k)[:n] result = FlashInferTrtllmDeferredFinalizeOutput( gemm2_out=gemm2_out, expert_weights=expert_weights,