[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

@@ -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,

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

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
         )

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."""

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}. "

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.")