From fc5a979f21218224b6b8852d661d5dcb3799d6dd Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 16 Sep 2026 01:50:29 -0700 Subject: [PATCH] [misc] Merge FlashInfer autotune caches across spec workers, pad MXFP4 TP shards, drop dead ngram attrs (#39678) Co-authored-by: BBuf <1182563586@qq.com> --- python/sglang/srt/layers/quantization/fp8.py | 2 +- .../srt/layers/quantization/fp8_utils.py | 3 +- .../mxfp4_flashinfer_trtllm_moe.py | 54 ++++++++++++++++++- python/sglang/srt/managers/scheduler.py | 5 -- .../model_runner_components/weight_updater.py | 2 +- .../runner/flashinfer_autotune.py | 14 +++-- 6 files changed, 68 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 92cf17af6..32af025f5 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -185,7 +185,7 @@ DSV4_DEQUANT_FP4_TABLE = torch.tensor( 3.0, 4.0, 6.0, - 0.0, + -0.0, -0.5, -1.0, -1.5, diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 361899487..27f732850 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1147,8 +1147,9 @@ def deepgemm_w8a8_block_fp8_linear_with_fallback( # TODO: https://github.com/sgl-project/sglang/pull/6890#issuecomment-2943395737 shape_supported = weight.shape[0] % 64 == 0 and weight.shape[1] % 128 == 0 + block_supported = list(block_size) == [128, 128] - if not (shape_supported and dtype_supported): + if not (shape_supported and dtype_supported and block_supported): # fall back to triton # If weight_scale is in UE8M0 packed format (int32), convert back to float32 # UE8M0 format has shape (N, K//block_k//4) with dtype int32 diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py index f178c5cc8..53d088690 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -22,7 +22,7 @@ from sglang.srt.utils import ( log_info_on_rank0, set_weight_attrs, ) -from sglang.srt.utils.common import next_power_of_2 +from sglang.srt.utils.common import next_power_of_2, print_warning_once _MXFP8_QUANTIZE_BACKEND = "cute-dsl" if get_platform().is_sm100 else "cuda" @@ -47,6 +47,56 @@ _USE_OFFICIAL_SHUFFLE = get_bool_env_var( ) +def _pad_intermediate_size(layer: Module) -> None: + intermediate_size = layer.w13_weight.shape[1] // 2 + padded_size = (intermediate_size + 127) // 128 * 128 + if padded_size == intermediate_size: + return + + # Gate and up occupy separate halves; each needs its own zero tail. + for name, fill_value in ( + ("w13_weight", 0), + ("w13_weight_scale_inv", 1), + ): + param = getattr(layer, name) + num_experts, _, width = param.shape + padded = torch.full( + (num_experts, 2 * padded_size, width), + fill_value, + dtype=param.dtype, + device=param.device, + ) + padded[:, :intermediate_size] = param[:, :intermediate_size] + padded[:, padded_size : padded_size + intermediate_size] = param[ + :, intermediate_size: + ] + param.data = padded + + for name, elements_per_column, fill_value in ( + ("w2_weight", 2, 0), + ("w2_weight_scale_inv", 32, 1), + ): + param = getattr(layer, name) + num_experts, hidden_size, width = param.shape + padded = torch.full( + (num_experts, hidden_size, padded_size // elements_per_column), + fill_value, + dtype=param.dtype, + device=param.device, + ) + padded[:, :, :width] = param + param.data = padded + + layer.intermediate_size_per_partition = padded_size + print_warning_once( + f"flashinfer_mxfp4 MoE padded the local intermediate size from " + f"{intermediate_size} to {padded_size} for 128-element kernel alignment " + "after TP weight loading. Padding adds unused channels and may waste " + "compute and memory. Use this TP MoE configuration with caution and " + "benchmark it against a TP/EP configuration that avoids padding." + ) + + class Mxfp4FlashinferTrtllmMoEMethod: fuse_routed_scaling_factor_in_topk = True @@ -150,6 +200,8 @@ class Mxfp4FlashinferTrtllmMoEMethod: if getattr(layer, "_mega_moe_weights_built", False): return + _pad_intermediate_size(layer) + w13_w, w13_s = reorder_w1w3_to_w3w1( layer.w13_weight.data, layer.w13_weight_scale_inv.data ) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index d63f4e8d6..b06d1829f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1680,11 +1680,6 @@ class Scheduler( self.tp_worker.model_runner.ngram_embedding_manager ) self.use_ngram_embedding = self.tp_worker.model_config.use_ngram_embedding - if self.use_ngram_embedding: - self.token_table = self.tp_worker.model_runner.ngram_embedding_manager.table - hf_config = self.tp_worker.model_config.hf_config - self.ngram_embedding_n = hf_config.ngram_embedding_n - self.ngram_embedding_k = hf_config.ngram_embedding_k def init_deterministic_inference_config(self): """Initialize deterministic inference configuration for different attention backends.""" diff --git a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py index e9dd1c702..43d358544 100644 --- a/python/sglang/srt/model_executor/model_runner_components/weight_updater.py +++ b/python/sglang/srt/model_executor/model_runner_components/weight_updater.py @@ -307,7 +307,7 @@ class WeightUpdater: ) reconstructed_tensors = bucket.reconstruct_tensors() self.get_model().load_weights(reconstructed_tensors) - return True, f"Succeeded to update parameter online." + return True, "Succeeded to update parameter online." except Exception as e: error_msg = ( f"Failed to update parameter online: {e}. " diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 51fa74ecd..94153a6cd 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py @@ -249,12 +249,13 @@ def _drop_diverged_autotune_cache( @contextlib.contextmanager def flashinfer_autotune_context(model_runner: ModelRunner, *, run_lm_head: bool): # The gate below decides on the same inputs load_configs does. - from flashinfer.autotuner import _collect_metadata, autotune + from flashinfer.autotuner import AutoTuner, _collect_metadata, autotune mr = model_runner cache_path = flashinfer_autotune_cache_path(mr) sync_group = _autotune_tactic_sync_group(mr.tp_group) - if envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get(): + reuse_cache = envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get() + if reuse_cache: autotune_cache = cache_path if sync_group is not None: _drop_diverged_autotune_cache(cache_path, sync_group, _collect_metadata()) @@ -277,16 +278,23 @@ def flashinfer_autotune_context(model_runner: ModelRunner, *, run_lm_head: bool) from sglang.srt.layers.logits_processor import autotune_dummy_run_mode skip_ops = get_flashinfer_autotune_skip_ops(mr) + # autotune(cache=...) clears all file-loaded tactics on entry, which would drop + # the target's tactics when the draft worker loads; load and save them by hand. + tuner = AutoTuner.get() + if reuse_cache and autotune_cache.is_file(): + tuner.load_configs(str(autotune_cache)) with ( _autotune_process_group(sync_group), autotune( True, - cache=str(autotune_cache), + cache=None if reuse_cache else str(autotune_cache), skip_ops=skip_ops, ), autotune_dummy_run_mode(run_lm_head=run_lm_head), ): yield + if reuse_cache: + tuner.save_configs(str(autotune_cache)) torch.cuda.current_stream().wait_stream(mr.forward_stream) logger.info("FlashInfer autotune completed.")