B300 NVFP4 port P1 stage-1: advance flashinfer_trtllm NVFP4 MoE + modelopt_fp4 path to upstream HEAD

Port the flashinfer-series NVFP4 path (GLM-5.2-NVFP4 + --moe-runner-backend flashinfer_trtllm
+ --fp4-gemm-backend flashinfer_trtllm) from upstream HEAD. b300-exp's flashinfer files were
byte-identical to fork-base (2d288ba8c9), so wholesale-replace == replaying the commit chain
without dragging baseline files forward (cherry-pick rejected: avg 17-54 files/commit entanglement).

Wholesale-replace (advanced to HEAD): modelopt_quant.py, moe_runner/flashinfer_trtllm.py,
flashinfer_trtllm_moe.py, fp4_utils.py, moe_runner/base.py. Brings split-w13 gate/up scales
(#27588), deferred-finalize symm-output (#27720), 4over6 + per-token activation (default-off).
ADD: marlin_utils_fp4.py, mxfp4_flashinfer_trtllm_moe.py (PackTopkIds), nvfp4_online.py (C1).
Surgical deps: environ +5 flags, common.alias_or_bind_derived_param, fp8_utils.apply_fp8_linear_bmm_flashinfer,
pynccl_allocator.is_tensor_in_symmetric_mempool, moe/utils.is_flashinfer_cutedsl_v1_path.

D1: gate modelopt-fp8 enable_flashinfer_bmm OFF by default (SGLANG_MODELOPT_FP8_FLASHINFER_BMM)
to keep the GLM-5.1-FP8 baseline byte-identical (upstream #28333 defaults it on for sm100).

Attention backend classes deferred (NSA bypasses them, verified 0.6.12-compatible).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/srt/distributed/device_communicators/pynccl_allocator.py

@@ -81,6 +81,21 @@ def is_symmetric_memory_enabled():
         return False
 
 
+def is_tensor_in_symmetric_mempool(tensor: torch.Tensor) -> bool:
+    """Check if a tensor's storage is allocated in the NCCL symmetric memory pool."""
+
+    if _mem_pool is None:
+        return False  # Pool not initialized
+
+    data_ptr = tensor.untyped_storage().data_ptr()
+
+    for segment in _mem_pool.snapshot():
+        for block in segment["blocks"]:
+            if block["address"] == data_ptr:
+                return True
+    return False
+
+
 def set_graph_pool_id(graph_pool_id):
     global _graph_pool_id
     _graph_pool_id = graph_pool_id

python/sglang/srt/environ.py

@@ -182,6 +182,15 @@ class Envs:
     SGLANG_SORT_WEIGHT_FILES = EnvBool(False)
     SGLANG_DISABLED_MODEL_ARCHS = EnvTuple(tuple())
 
+    # NVFP4 / flashinfer (B300 port)
+    SGLANG_EXPERIMENTAL_LORA_OPTI = EnvBool(False)
+    SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION = EnvBool(False)
+    FLASHINFER_NVFP4_4OVER6 = EnvBool(False)
+    FLASHINFER_NVFP4_4OVER6_E4M3_USE_256 = EnvBool(False)
+    # B300 port D1: gate modelopt-fp8 flashinfer-bmm OFF by default to keep the
+    # GLM-5.1-FP8 baseline byte-identical (upstream #28333 defaults it on for sm100).
+    SGLANG_MODELOPT_FP8_FLASHINFER_BMM = EnvBool(False)
+
     # Logging Options
     SGLANG_LOG_GC = EnvBool(False)
     SGLANG_LOG_FORWARD_ITERS = EnvBool(False)

python/sglang/srt/layers/moe/flashinfer_trtllm_moe.py

@@ -28,6 +28,7 @@ def _fake_fp8_block_scale_moe(
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
     fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     return torch.empty(
         hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
@@ -58,6 +59,7 @@ def trtllm_fp8_block_scale_moe_wrapper(
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
     fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     try:
         from flashinfer.fused_moe import trtllm_fp8_block_scale_moe
@@ -94,6 +96,11 @@ def trtllm_fp8_block_scale_moe_wrapper(
 
         kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type)
 
+    if activation_type is not None:
+        from flashinfer.fused_moe.core import ActivationType
+
+        kwargs["activation_type"] = ActivationType(activation_type)
+
     return trtllm_fp8_block_scale_moe(**kwargs)
 
 
@@ -120,6 +127,7 @@ def _fake_fp8_block_scale_routed_moe(
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
     fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     return torch.empty(
         hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
@@ -150,6 +158,7 @@ def trtllm_fp8_block_scale_routed_moe_wrapper(
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
     fp8_quantization_type: Optional[int] = None,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     try:
         from flashinfer.fused_moe import trtllm_fp8_block_scale_routed_moe
@@ -186,6 +195,11 @@ def trtllm_fp8_block_scale_routed_moe_wrapper(
 
         kwargs["fp8_quantization_type"] = Fp8QuantizationType(fp8_quantization_type)
 
+    if activation_type is not None:
+        from flashinfer.fused_moe.core import ActivationType
+
+        kwargs["activation_type"] = ActivationType(activation_type)
+
     return trtllm_fp8_block_scale_routed_moe(**kwargs)
 
 
@@ -210,6 +224,7 @@ def _fake_fp8_per_tensor_scale_moe(
     routing_method_type: int = 0,
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     return torch.empty(
         hidden_states.shape, dtype=torch.bfloat16, device=hidden_states.device
@@ -238,6 +253,7 @@ def trtllm_fp8_per_tensor_scale_moe_wrapper(
     routing_method_type: int = 0,
     enable_pdl: Optional[bool] = None,
     tune_max_num_tokens: int = 8192,
+    activation_type: Optional[int] = None,
 ) -> torch.Tensor:
     # lazy import
     try:
@@ -271,4 +287,9 @@ def trtllm_fp8_per_tensor_scale_moe_wrapper(
         "tune_max_num_tokens": tune_max_num_tokens,
     }
 
+    if activation_type is not None:
+        from flashinfer.fused_moe.core import ActivationType
+
+        kwargs["activation_type"] = ActivationType(activation_type)
+
     return trtllm_fp8_per_tensor_scale_moe(**kwargs)

python/sglang/srt/layers/moe/moe_runner/base.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import contextvars
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Callable, Optional, Tuple, TypeGuard
+from typing import TYPE_CHECKING, Any, Callable, Generator, Optional, Tuple, TypeGuard
 
 import torch
 
@@ -26,6 +28,20 @@ if TYPE_CHECKING:
     )
 
 
+_moe_output_buf: contextvars.ContextVar[Optional[torch.Tensor]] = (
+    contextvars.ContextVar("moe_output_buf", default=None)
+)
+
+
+@contextmanager
+def moe_output_buffer_ctx(buf: torch.Tensor) -> Generator[None, None, None]:
+    token = _moe_output_buf.set(buf)
+    try:
+        yield
+    finally:
+        _moe_output_buf.reset(token)
+
+
 @dataclass
 class MoeRunnerConfig:
     # MoE parameters
@@ -48,6 +64,7 @@ class MoeRunnerConfig:
     routed_scaling_factor: Optional[float] = None
     gemm1_alpha: Optional[float] = None
     gemm1_clamp_limit: Optional[float] = None
+    swiglu_limit: Optional[float] = None
 
 
 @dataclass
@@ -82,7 +99,11 @@ class MoeRunnerCore(ABC):
 
     @abstractmethod
     def run(
-        self, runner_input: RunnerInput, quant_info: MoeQuantInfo, running_state: dict
+        self,
+        runner_input: RunnerInput,
+        quant_info: MoeQuantInfo,
+        running_state: dict,
+        hooks: Optional[Any] = None,
     ) -> RunnerOutput:
         pass
 

python/sglang/srt/layers/moe/utils.py

@@ -200,6 +200,14 @@ def get_moe_runner_backend() -> MoeRunnerBackend:
     return MOE_RUNNER_BACKEND
 
 
+def is_flashinfer_cutedsl_v1_path() -> bool:
+    """CuteDSL v1 + DeepEP low-latency path (no MoeRunner, no autotune)."""
+    return (
+        get_moe_runner_backend().is_flashinfer_cutedsl()
+        and get_moe_a2a_backend().is_deepep()
+    )
+
+
 def get_speculative_moe_runner_backend() -> MoeRunnerBackend:
     global SPECULATIVE_MOE_RUNNER_BACKEND
     if SPECULATIVE_MOE_RUNNER_BACKEND is None:

python/sglang/srt/layers/quantization/fp4_utils.py

@@ -2,10 +2,16 @@ from __future__ import annotations
 
 import logging
 from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
-from sglang.srt.environ import envs
-from sglang.srt.utils.common import is_sm120_supported
+import torch
+
+from sglang.srt.utils.common import (
+    get_device_capability,
+    is_cuda,
+    is_sm100_supported,
+)
+from sglang.srt.utils.custom_op import register_custom_op_from_extern
 
 if TYPE_CHECKING:
     from sglang.srt.server_args import ServerArgs
@@ -13,17 +19,94 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+fp4_quantize = None
+try:
+    from flashinfer import fp4_quantize as _flashinfer_fp4_quantize
+
+    _flashinfer_fp4_quantize_backend = "cute-dsl" if is_sm100_supported() else "cuda"
+
+    def _round_up(x: int, y: int) -> int:
+        return ((x + y - 1) // y) * y
+
+    def _flashinfer_fp4_quantize_impl(
+        input: torch.Tensor,
+        global_scale: Optional[torch.Tensor] = None,
+        sf_vec_size: int = 16,
+        sf_use_ue8m0: bool = False,
+        is_sf_swizzled_layout: bool = True,
+        is_sf_8x4_layout: bool = False,
+        enable_pdl: Optional[bool] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return _flashinfer_fp4_quantize(
+            input=input,
+            global_scale=global_scale,
+            sf_vec_size=sf_vec_size,
+            sf_use_ue8m0=sf_use_ue8m0,
+            is_sf_swizzled_layout=is_sf_swizzled_layout,
+            is_sf_8x4_layout=is_sf_8x4_layout,
+            enable_pdl=enable_pdl,
+            backend=_flashinfer_fp4_quantize_backend,
+        )
+
+    def _flashinfer_fp4_quantize_fake(
+        input: torch.Tensor,
+        global_scale: Optional[torch.Tensor] = None,
+        sf_vec_size: int = 16,
+        sf_use_ue8m0: bool = False,
+        is_sf_swizzled_layout: bool = True,
+        is_sf_8x4_layout: bool = False,
+        enable_pdl: Optional[bool] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        is_column_major = input.stride(-2) == 1
+        if is_column_major:
+            m = input.shape[-1]
+            K = input.shape[-2]
+        else:
+            m = input.numel() // input.shape[-1]
+            K = input.shape[-1]
+        if is_column_major:
+            x_q = input.new_empty((*input.shape[:-2], K // 2, m), dtype=torch.uint8)
+        else:
+            x_q = input.new_empty((*input.shape[:-1], K // 2), dtype=torch.uint8)
+        if is_sf_swizzled_layout:
+            row_size = 8 if is_sf_8x4_layout else 128
+            sf_rows = _round_up(m, row_size)
+            sf_cols = _round_up(K // sf_vec_size, 4)
+        else:
+            sf_rows = m
+            sf_cols = K // sf_vec_size
+        if is_column_major:
+            sf = input.new_empty((sf_cols, sf_rows), dtype=torch.uint8)
+        else:
+            sf = input.new_empty((sf_rows, sf_cols), dtype=torch.uint8)
+        return x_q, sf
+
+    fp4_quantize = register_custom_op_from_extern(
+        _flashinfer_fp4_quantize_impl,
+        op_name="flashinfer_fp4_quantize",
+        fake_impl=_flashinfer_fp4_quantize_fake,
+    )
+except ImportError:
+    fp4_quantize = None
+
+
 class Fp4GemmRunnerBackend(Enum):
     """Enum for FP4 GEMM runner backend selection."""
 
     AUTO = "auto"
+    CUTLASS = "cutlass"
     FLASHINFER_CUDNN = "flashinfer_cudnn"
+    FLASHINFER_CUTEDSL = "flashinfer_cutedsl"
     FLASHINFER_CUTLASS = "flashinfer_cutlass"
     FLASHINFER_TRTLLM = "flashinfer_trtllm"
+    MARLIN = "marlin"
 
     def is_auto(self) -> bool:
         return self == Fp4GemmRunnerBackend.AUTO
 
+    def is_cutlass(self) -> bool:
+        return self == Fp4GemmRunnerBackend.CUTLASS
+
     def is_flashinfer_cudnn(self) -> bool:
         return self == Fp4GemmRunnerBackend.FLASHINFER_CUDNN
 
@@ -33,6 +116,15 @@ class Fp4GemmRunnerBackend(Enum):
     def is_flashinfer_trtllm(self) -> bool:
         return self == Fp4GemmRunnerBackend.FLASHINFER_TRTLLM
 
+    def is_flashinfer_cutedsl(self) -> bool:
+        return self == Fp4GemmRunnerBackend.FLASHINFER_CUTEDSL
+
+    def is_marlin(self) -> bool:
+        return self == Fp4GemmRunnerBackend.MARLIN
+
+    def is_flashinfer(self) -> bool:
+        return self.value.startswith("flashinfer_")
+
     def get_flashinfer_backend(self) -> str:
         """Get the backend string to pass to FlashInfer's mm_fp4 API.
 
@@ -41,7 +133,10 @@ class Fp4GemmRunnerBackend(Enum):
             'flashinfer_trtllm' -> 'trtllm'
             'flashinfer_cutlass' -> 'cutlass'
             'flashinfer_cudnn' -> 'cudnn'
+            'flashinfer_cutedsl' -> 'cute-dsl'
         """
+        if self == Fp4GemmRunnerBackend.FLASHINFER_CUTEDSL:
+            return "cute-dsl"
         if self.value.startswith("flashinfer_"):
             return self.value.removeprefix("flashinfer_")
         else:
@@ -56,36 +151,11 @@ def initialize_fp4_gemm_config(server_args: ServerArgs) -> None:
     global FP4_GEMM_RUNNER_BACKEND
 
     backend = server_args.fp4_gemm_runner_backend
-
-    # Handle deprecated env var for backward compatibility
-    # TODO: Remove this in a future version
-    if envs.SGLANG_FLASHINFER_FP4_GEMM_BACKEND.is_set():
-        env_backend = envs.SGLANG_FLASHINFER_FP4_GEMM_BACKEND.get()
-        if backend == "auto":
-            logger.warning(
-                "SGLANG_FLASHINFER_FP4_GEMM_BACKEND is deprecated. "
-                f"Please use '--fp4-gemm-backend={env_backend}' instead."
-            )
-            if not env_backend.startswith("flashinfer_"):
-                env_backend = "flashinfer_" + env_backend
-            backend = env_backend
-        else:
-            logger.warning(
-                f"FP4 GEMM backend set to '{backend}' via --fp4-gemm-backend overrides "
-                "environment variable SGLANG_FLASHINFER_FP4_GEMM_BACKEND. "
-                "Using server argument value."
-            )
-
     if backend == "auto":
-        if is_sm120_supported():
-            # flashinfer_cutlass produces NaN in dense MLP layers with
-            # heterogeneous batches on SM120 (Blackwell).  cudnn is stable.
-            # See: https://github.com/sgl-project/sglang/issues/20043
-            backend = "flashinfer_cudnn"
-            logger.info(
-                "SM120 (Blackwell) detected: auto-selecting "
-                "fp4-gemm-backend=flashinfer_cudnn"
-            )
+        if is_sm100_supported():
+            backend = "flashinfer_cutedsl"
+        elif is_cuda() and (10, 0) > get_device_capability() >= (8, 0):
+            backend = "marlin"
         else:
             backend = "flashinfer_cutlass"
 

python/sglang/srt/layers/quantization/fp8_utils.py

@@ -211,6 +211,7 @@ def _check_cutlass_block_fp8_hardware_support() -> bool:
 
 
 if is_blackwell_supported() and is_flashinfer_available():
+    from flashinfer import bmm_fp8 as flashinfer_bmm_fp8
     from flashinfer import mm_mxfp8 as _raw_flashinfer_mm_mxfp8
     from flashinfer import mxfp8_quantize as _raw_flashinfer_mxfp8_quantize
     from flashinfer.gemm import gemm_fp8_nt_groupwise as _raw_gemm_fp8_nt_groupwise
@@ -1368,6 +1369,28 @@ def _apply_fallback_scaled_mm(
     return output.to(dtype=input_dtype)
 
 
+def apply_fp8_linear_bmm_flashinfer(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    input_scale: torch.Tensor,
+    bias: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Per-tensor static fp8 linear via flashinfer bmm_fp8 (SM10X only).
+
+    B300 port: gated OFF by default (SGLANG_MODELOPT_FP8_FLASHINFER_BMM). `static_quant_fp8`
+    is a module-level import; `flashinfer_bmm_fp8` is imported under the Blackwell guard above
+    (this path only runs on sm100, where that guard is taken).
+    """
+    output_shape = [*input.shape[:-1], weight.shape[1]]
+    input_2d = input.view(-1, input.shape[-1])
+    qinput, x_scale = static_quant_fp8(input_2d, input_scale, repeat_scale=False)
+    output = flashinfer_bmm_fp8(qinput, weight, x_scale, weight_scale, input.dtype)
+    if bias is not None:
+        output = output + bias
+    return output.view(*output_shape)
+
+
 def apply_fp8_linear(
     input: torch.Tensor,
     weight: torch.Tensor,

python/sglang/srt/layers/quantization/marlin_utils_fp4.py

@@ -0,0 +1,518 @@
+from __future__ import annotations
+
+import torch
+
+from sglang.srt.layers.quantization.marlin_utils import (
+    USE_FP32_REDUCE_DEFAULT,
+    marlin_make_workspace,
+    marlin_permute_bias,
+    marlin_permute_scales,
+    should_use_atomic_add_reduce,
+)
+from sglang.srt.layers.quantization.utils import get_scalar_types
+from sglang.srt.utils import is_cuda
+from sglang.srt.utils.custom_op import register_custom_op
+
+_is_cuda = is_cuda()
+
+if _is_cuda:
+    from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm
+    from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
+
+ScalarType, scalar_types = get_scalar_types()
+
+
+def nvfp4_marlin_process_scales(marlin_scales: torch.Tensor) -> torch.Tensor:
+    if not (marlin_scales >= 0).all():
+        # NVFP4 ModelOpt scales are expected to be non-negative. Keep this as
+        # a warning so unusual checkpoints can still load for diagnosis.
+        import logging
+
+        logging.getLogger(__name__).warning_once(
+            "NVFP4 Marlin assumes non-negative scales, but negative scales "
+            "were found. Accuracy may be degraded."
+        )
+
+    marlin_scales = marlin_scales.to(torch.half)
+    marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(
+        marlin_scales.size(0), -1
+    )
+    marlin_scales = (marlin_scales * (2**7)).view(torch.int16) << 1
+    marlin_scales = marlin_scales.view(torch.float8_e4m3fn)
+    return marlin_scales[:, 1::2].contiguous()
+
+
+def nvfp4_marlin_process_global_scale(global_scale: torch.Tensor) -> torch.Tensor:
+    assert global_scale.dtype in [torch.half, torch.bfloat16]
+    global_scale_shape = global_scale.shape
+    fp4_exponent = 2
+    if global_scale.dtype == torch.half:
+        target_exponent = 5
+    elif global_scale.dtype == torch.bfloat16:
+        target_exponent = 8
+    exponent_bias = 2 ** (target_exponent - 1) - 2 ** (fp4_exponent - 1)
+    global_scale = global_scale * (2.0 ** (exponent_bias - 7))
+    if global_scale_shape == torch.Size([]):
+        global_scale = global_scale.reshape(1)
+    return global_scale
+
+
+def fake_apply_fp4_marlin_linear(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_global_scale: torch.Tensor,
+    workspace: torch.Tensor,
+    size_n: int,
+    size_k: int,
+    bias: torch.Tensor | None = None,
+    use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,
+) -> torch.Tensor:
+    del weight, weight_scale, weight_global_scale, workspace, size_k, bias
+    out_shape = input.shape[:-1] + (size_n,)
+    return input.new_empty(out_shape)
+
+
+@register_custom_op(fake_impl=fake_apply_fp4_marlin_linear)
+def apply_fp4_marlin_linear(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    weight_global_scale: torch.Tensor,
+    workspace: torch.Tensor,
+    size_n: int,
+    size_k: int,
+    bias: torch.Tensor | None = None,
+    use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT,
+) -> torch.Tensor:
+    if input.dtype not in (torch.float16, torch.bfloat16):
+        raise RuntimeError("NVFP4 Marlin requires FP16 or BF16 activations.")
+
+    reshaped_x = input.reshape(-1, input.shape[-1])
+    out_shape = input.shape[:-1] + (size_n,)
+
+    use_atomic_add = should_use_atomic_add_reduce(
+        m=reshaped_x.size(0),
+        n=size_n,
+        k=size_k,
+        device=input.device,
+        dtype=input.dtype,
+    )
+
+    output = gptq_marlin_gemm(
+        a=reshaped_x,
+        c=None,
+        b_q_weight=weight,
+        b_scales=weight_scale,
+        global_scale=weight_global_scale,
+        b_zeros=None,
+        g_idx=None,
+        perm=None,
+        workspace=workspace,
+        b_q_type=scalar_types.float4_e2m1f,
+        size_m=reshaped_x.size(0),
+        size_n=size_n,
+        size_k=size_k,
+        is_k_full=True,
+        use_atomic_add=use_atomic_add,
+        use_fp32_reduce=use_fp32_reduce,
+    )
+
+    if bias is not None:
+        output.add_(bias)
+
+    return output.reshape(out_shape)
+
+
+def prepare_nvfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
+    if getattr(layer, "quant_config", None) is not None:
+        group_size = layer.quant_config.group_size
+        if group_size != 16:
+            raise ValueError(f"NVFP4 Marlin requires group_size=16, got {group_size}.")
+
+    part_size_n = layer.output_size_per_partition
+    part_size_k = layer.input_size_per_partition
+    param_dtype = getattr(layer, "params_dtype", getattr(layer, "orig_dtype", None))
+    if param_dtype not in (torch.float16, torch.bfloat16):
+        raise RuntimeError("NVFP4 Marlin requires FP16 or BF16 activation dtype.")
+
+    assert layer.weight.shape == (part_size_n, part_size_k // 2)
+
+    if part_size_n % 64 != 0:
+        raise ValueError(
+            f"NVFP4 Marlin requires output_size_per_partition to be a multiple of 64, "
+            f"got {part_size_n}."
+        )
+
+    device = layer.weight.device
+    layer.workspace = marlin_make_workspace(device)
+
+    perm = torch.empty(0, dtype=torch.int, device=device)
+    qweight = layer.weight.view(torch.int32).T.contiguous()
+    marlin_qweight = gptq_marlin_repack(
+        b_q_weight=qweight,
+        perm=perm,
+        size_k=part_size_k,
+        size_n=part_size_n,
+        num_bits=4,
+    )
+    layer.weight = torch.nn.Parameter(marlin_qweight, requires_grad=False)
+
+    weight_scale = layer.weight_scale.T.contiguous().to(param_dtype)
+    weight_scale = marlin_permute_scales(
+        s=weight_scale,
+        size_k=part_size_k,
+        size_n=part_size_n,
+        group_size=16,
+    )
+    weight_scale = nvfp4_marlin_process_scales(weight_scale)
+    layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+
+    weight_global_scale = layer.weight_global_scale.to(param_dtype)
+    weight_global_scale = nvfp4_marlin_process_global_scale(weight_global_scale)
+    layer.weight_global_scale = torch.nn.Parameter(
+        weight_global_scale, requires_grad=False
+    )
+
+    if hasattr(layer, "bias") and layer.bias is not None:
+        assert layer.bias.shape == (part_size_n,)
+        bias = marlin_permute_bias(layer.bias)
+        layer.bias = torch.nn.Parameter(bias, requires_grad=False)
+
+
+def mxfp4_marlin_process_scales(
+    marlin_scales: torch.Tensor,
+    input_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    if input_dtype is None or input_dtype.itemsize == 2:
+        marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view(
+            marlin_scales.size(0), -1
+        )
+    marlin_scales = marlin_scales.to(torch.float8_e8m0fnu)
+    if input_dtype == torch.float8_e4m3fn:
+        marlin_scales = marlin_scales.view(torch.uint8)
+        assert marlin_scales.max() <= 249
+        # exponent_bias (fp4->fp8) = 2 ** 3 - 2 ** 1 = 6
+        marlin_scales = marlin_scales + 6
+        marlin_scales = marlin_scales.view(torch.float8_e8m0fnu)
+    return marlin_scales
+
+
+def _normalize_scale_tensor(
+    scales: torch.Tensor, target_dtype: torch.dtype
+) -> torch.Tensor:
+    # The kernel consumes E8M0 exponents. Regardless of the placeholder dtype
+    # the loader used, we want the *numerical* value 2**e in ``target_dtype``.
+    # float32/bfloat16/float16 containers hold the numerical 2**e directly
+    # (they were filled via a dtype-promoting copy from uint8/e8m0).
+    # uint8/int8 containers hold the raw E8M0 byte and must be reinterpreted.
+    if scales.dtype == torch.float8_e8m0fnu:
+        return scales.to(target_dtype)
+    if scales.dtype == torch.uint8:
+        return scales.view(torch.float8_e8m0fnu).to(target_dtype)
+    if scales.dtype == torch.int8:
+        return scales.view(torch.uint8).view(torch.float8_e8m0fnu).to(target_dtype)
+    if scales.dtype in (torch.float32, torch.bfloat16, torch.float16):
+        return scales.to(target_dtype)
+    raise TypeError(f"Unsupported MXFP4 scale dtype for Marlin: {scales.dtype}")
+
+
+def _get_optional_param(layer: torch.nn.Module, *names: str) -> torch.Tensor | None:
+    for name in names:
+        value = getattr(layer, name, None)
+        if value is not None:
+            return value
+    return None
+
+
+def deinterleave_moe_mxfp4_w13_for_marlin(layer: torch.nn.Module) -> None:
+    """Convert GPT-OSS interleaved w13 rows to Marlin's contiguous halves.
+
+    GPT-OSS stores gate/up rows as [gate0, up0, gate1, up1, ...]. The Marlin
+    fused activation consumes [all_gate_rows, all_up_rows].
+    """
+
+    w13 = layer.w13_weight.data
+    w13_scale = _get_optional_param(layer, "w13_weight_scale", "w13_weight_scale_inv")
+    w13_bias = _get_optional_param(layer, "w13_weight_bias", "w13_bias")
+
+    if w13.shape[1] % 2 != 0:
+        raise ValueError(f"Expected even w13 row dimension, got {w13.shape}.")
+
+    e, n, k = w13.shape
+    layer.w13_weight.data = (
+        w13.view(e, n // 2, 2, k).permute(0, 2, 1, 3).contiguous().view(e, n, k)
+    )
+
+    if w13_scale is not None:
+        scale = w13_scale.data
+        if scale.shape[1] != n:
+            raise ValueError(
+                f"Expected w13 scale row dimension {n}, got {scale.shape}."
+            )
+        w13_scale.data = (
+            scale.view(e, n // 2, 2, scale.shape[-1])
+            .permute(0, 2, 1, 3)
+            .contiguous()
+            .view(e, n, scale.shape[-1])
+        )
+
+    if w13_bias is not None:
+        bias = w13_bias.data
+        if bias.shape[1] != n:
+            raise ValueError(f"Expected w13 bias row dimension {n}, got {bias.shape}.")
+        w13_bias.data = bias.view(e, n // 2, 2).permute(0, 2, 1).contiguous().view(e, n)
+
+
+def prepare_moe_mxfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
+    group_size = 32
+    w13 = layer.w13_weight.data
+    w2 = layer.w2_weight.data
+    w13_scale = _get_optional_param(layer, "w13_weight_scale", "w13_weight_scale_inv")
+    w2_scale = _get_optional_param(layer, "w2_weight_scale", "w2_weight_scale_inv")
+    w13_bias = _get_optional_param(layer, "w13_weight_bias", "w13_bias")
+    w2_bias = _get_optional_param(layer, "w2_weight_bias", "w2_bias")
+
+    if w13_scale is None or w2_scale is None:
+        raise ValueError("MXFP4 Marlin requires w13/w2 weight scales.")
+
+    w13_scale_data = w13_scale.data if hasattr(w13_scale, "data") else w13_scale
+    w2_scale_data = w2_scale.data if hasattr(w2_scale, "data") else w2_scale
+    w13_bias_data = w13_bias.data if hasattr(w13_bias, "data") else w13_bias
+    w2_bias_data = w2_bias.data if hasattr(w2_bias, "data") else w2_bias
+
+    num_experts = w13.shape[0]
+    intermediate_size = w13.shape[1] // 2
+    hidden_size = w13.shape[2] * 2
+    if hidden_size % 128 == 0:
+        padded_intermediate_size = ((intermediate_size + 63) // 64) * 64
+    else:
+        if hidden_size % 64 != 0:
+            raise ValueError(
+                f"MXFP4 Marlin requires hidden_size to be divisible by 64, "
+                f"got {hidden_size}."
+            )
+        padded_intermediate_size = ((intermediate_size + 127) // 128) * 128
+    param_dtype = getattr(
+        layer,
+        "orig_dtype",
+        w13_bias_data.dtype if w13_bias_data is not None else torch.bfloat16,
+    )
+
+    device = w13.device
+    layer.workspace = marlin_make_workspace(device, 4)
+    perm = torch.empty(0, dtype=torch.int, device=device)
+
+    def _pad_w13(x: torch.Tensor) -> torch.Tensor:
+        if padded_intermediate_size == intermediate_size:
+            return x
+        x = x.view(num_experts, 2, intermediate_size, x.shape[-1])
+        x = torch.nn.functional.pad(
+            x, (0, 0, 0, padded_intermediate_size - intermediate_size)
+        )
+        return x.reshape(num_experts, 2 * padded_intermediate_size, -1)
+
+    def _pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor:
+        if padded_intermediate_size == intermediate_size:
+            return x
+        return torch.nn.functional.pad(
+            x, (0, (padded_intermediate_size - intermediate_size) // packing)
+        )
+
+    w13 = _pad_w13(w13)
+    w2 = _pad_w2(w2, packing=2)
+    w13_scale_data = _pad_w13(_normalize_scale_tensor(w13_scale_data, param_dtype))
+    w2_scale_data = _pad_w2(
+        _normalize_scale_tensor(w2_scale_data, param_dtype),
+        packing=group_size,
+    )
+    if w13_bias_data is not None:
+        w13_bias_data = _pad_w13(w13_bias_data.unsqueeze(-1)).squeeze(-1)
+
+    def _repack_weight(weight: torch.Tensor, is_w13: bool) -> torch.Tensor:
+        if is_w13:
+            size_n, size_k = padded_intermediate_size * 2, hidden_size
+        else:
+            size_n, size_k = hidden_size, padded_intermediate_size
+        assert weight.shape == (num_experts, size_n, size_k // 2)
+
+        tensor_list = []
+        for i in range(num_experts):
+            qweight = weight[i].view(torch.int32).T.contiguous()
+            marlin_qweight = gptq_marlin_repack(
+                b_q_weight=qweight,
+                perm=perm,
+                size_k=size_k,
+                size_n=size_n,
+                num_bits=4,
+            )
+            tensor_list.append(marlin_qweight)
+        return torch.stack(tensor_list)
+
+    def _permute_scales(scales: torch.Tensor, is_w13: bool) -> torch.Tensor:
+        if is_w13:
+            size_n, size_k = padded_intermediate_size * 2, hidden_size
+        else:
+            size_n, size_k = hidden_size, padded_intermediate_size
+
+        tensor_list = []
+        for i in range(num_experts):
+            scale = scales[i].T.contiguous()
+            marlin_scales = marlin_permute_scales(
+                s=scale,
+                size_k=size_k,
+                size_n=size_n,
+                group_size=group_size,
+            )
+            tensor_list.append(
+                mxfp4_marlin_process_scales(
+                    marlin_scales,
+                    input_dtype=param_dtype,
+                )
+            )
+        return torch.stack(tensor_list)
+
+    def _permute_bias(bias: torch.Tensor | None) -> torch.Tensor | None:
+        if bias is None:
+            return None
+        tensor_list = []
+        for i in range(num_experts):
+            tensor_list.append(marlin_permute_bias(bias[i].to(param_dtype)))
+        return torch.stack(tensor_list)
+
+    w13_marlin = _repack_weight(w13, True)
+    w2_marlin = _repack_weight(w2, False)
+    w13_scale_marlin = _permute_scales(w13_scale_data, True)
+    w2_scale_marlin = _permute_scales(w2_scale_data, False)
+
+    layer.w13_weight = torch.nn.Parameter(w13_marlin, requires_grad=False)
+    layer.w2_weight = torch.nn.Parameter(w2_marlin, requires_grad=False)
+    layer.w13_weight_scale = torch.nn.Parameter(w13_scale_marlin, requires_grad=False)
+    layer.w2_weight_scale = torch.nn.Parameter(w2_scale_marlin, requires_grad=False)
+
+    if w13_bias_data is not None:
+        layer.w13_weight_bias = torch.nn.Parameter(
+            _permute_bias(w13_bias_data), requires_grad=False
+        )
+    if w2_bias_data is not None:
+        layer.w2_weight_bias = torch.nn.Parameter(
+            _permute_bias(w2_bias_data), requires_grad=False
+        )
+
+
+def prepare_moe_nvfp4_layer_for_marlin(layer: torch.nn.Module) -> None:
+    if layer.quant_config.group_size != 16:
+        raise ValueError(
+            f"NVFP4 Marlin MoE requires group_size=16, got {layer.quant_config.group_size}."
+        )
+
+    w13 = layer.w13_weight.data
+    w2 = layer.w2_weight.data
+    w13_scale = layer.w13_weight_scale.data
+    w2_scale = layer.w2_weight_scale.data
+    w13_global_scale = layer.w13_weight_scale_2.data
+    w2_global_scale = layer.w2_weight_scale_2.data
+    w13_bias = getattr(layer, "w13_bias", None)
+    w2_bias = getattr(layer, "w2_bias", None)
+
+    num_experts = w13.shape[0]
+    num_shards = 2 if layer.moe_runner_config.is_gated else 1
+    intermediate_size = layer.intermediate_size_per_partition
+    hidden_size = w13.shape[2] * 2
+    param_dtype = layer.params_dtype
+    if param_dtype not in (torch.float16, torch.bfloat16):
+        raise RuntimeError("NVFP4 Marlin MoE requires FP16 or BF16 activations.")
+
+    device = w13.device
+    layer.workspace = marlin_make_workspace(device, 4)
+    perm = torch.empty(0, dtype=torch.int, device=device)
+
+    if not layer.moe_runner_config.is_gated:
+        padded_intermediate_size = ((intermediate_size + 127) // 128) * 128
+        intermediate_size_pad = padded_intermediate_size - intermediate_size
+        if intermediate_size_pad:
+            w13 = torch.nn.functional.pad(w13, (0, 0, 0, intermediate_size_pad))
+            w13_scale = torch.nn.functional.pad(
+                w13_scale, (0, 0, 0, intermediate_size_pad)
+            )
+            w2 = torch.nn.functional.pad(w2, (0, intermediate_size_pad // 2, 0, 0))
+            w2_scale = torch.nn.functional.pad(
+                w2_scale, (0, intermediate_size_pad // 16)
+            )
+            if w13_bias is not None:
+                w13_bias = torch.nn.functional.pad(w13_bias, (0, intermediate_size_pad))
+            intermediate_size = padded_intermediate_size
+
+    def _repack_weight(weight: torch.Tensor, is_w13: bool) -> torch.Tensor:
+        if is_w13:
+            size_n, size_k = intermediate_size * num_shards, hidden_size
+        else:
+            size_n, size_k = hidden_size, intermediate_size
+        assert weight.shape == (num_experts, size_n, size_k // 2)
+
+        tensor_list = []
+        for i in range(num_experts):
+            qweight = weight[i].view(torch.int32).T.contiguous()
+            marlin_qweight = gptq_marlin_repack(
+                b_q_weight=qweight,
+                perm=perm,
+                size_k=size_k,
+                size_n=size_n,
+                num_bits=4,
+            )
+            tensor_list.append(marlin_qweight)
+        return torch.stack(tensor_list)
+
+    def _permute_scales(scales: torch.Tensor, is_w13: bool) -> torch.Tensor:
+        scales = scales.to(param_dtype)
+        if is_w13:
+            size_n, size_k = intermediate_size * num_shards, hidden_size
+        else:
+            size_n, size_k = hidden_size, intermediate_size
+
+        tensor_list = []
+        for i in range(num_experts):
+            scale = scales[i].T.contiguous()
+            marlin_scales = marlin_permute_scales(
+                s=scale,
+                size_k=size_k,
+                size_n=size_n,
+                group_size=16,
+            )
+            tensor_list.append(nvfp4_marlin_process_scales(marlin_scales))
+        return torch.stack(tensor_list)
+
+    def _process_global_scale(global_scale: torch.Tensor) -> torch.Tensor:
+        return nvfp4_marlin_process_global_scale(global_scale.to(param_dtype))
+
+    def _permute_bias(bias: torch.Tensor | None) -> torch.Tensor | None:
+        if bias is None:
+            return None
+        tensor_list = []
+        for i in range(num_experts):
+            tensor_list.append(marlin_permute_bias(bias[i].to(param_dtype)))
+        return torch.stack(tensor_list)
+
+    layer.w13_weight = torch.nn.Parameter(
+        _repack_weight(w13, True), requires_grad=False
+    )
+    layer.w2_weight = torch.nn.Parameter(_repack_weight(w2, False), requires_grad=False)
+    layer.w13_weight_scale = torch.nn.Parameter(
+        _permute_scales(w13_scale, True), requires_grad=False
+    )
+    layer.w2_weight_scale = torch.nn.Parameter(
+        _permute_scales(w2_scale, False), requires_grad=False
+    )
+    layer.w13_weight_scale_2 = torch.nn.Parameter(
+        _process_global_scale(w13_global_scale), requires_grad=False
+    )
+    layer.w2_weight_scale_2 = torch.nn.Parameter(
+        _process_global_scale(w2_global_scale), requires_grad=False
+    )
+
+    if w13_bias is not None:
+        layer.w13_bias = torch.nn.Parameter(
+            _permute_bias(w13_bias), requires_grad=False
+        )
+    if w2_bias is not None:
+        layer.w2_bias = torch.nn.Parameter(_permute_bias(w2_bias), requires_grad=False)

python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py

@@ -0,0 +1,474 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import torch
+import triton
+import triton.language as tl
+from torch.nn import Module
+from torch.nn.parameter import Parameter
+
+from sglang.srt.distributed import get_tp_group
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
+from sglang.srt.layers.dp_attention import is_allocation_symmetric
+from sglang.srt.layers.moe.utils import RoutingMethodType
+from sglang.srt.server_args import get_global_server_args
+from sglang.srt.utils import (
+    is_flashinfer_available,
+    log_info_on_rank0,
+    set_weight_attrs,
+)
+from sglang.srt.utils.common import is_sm100_supported, next_power_of_2
+
+_MXFP8_QUANTIZE_BACKEND = "cute-dsl" if is_sm100_supported() else "cuda"
+
+if is_flashinfer_available():
+    from flashinfer import mxfp8_quantize, shuffle_matrix_a, shuffle_matrix_sf_a
+    from flashinfer.fp4_quantization import block_scale_interleave
+    from flashinfer.fused_moe import trtllm_fp4_block_scale_routed_moe
+    from flashinfer.fused_moe.core import (
+        _maybe_get_cached_w3_w1_permute_indices,
+        get_w2_permute_indices_with_cache,
+    )
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput
+
+from sglang.srt.utils.common import get_bool_env_var
+
+_USE_OFFICIAL_SHUFFLE = get_bool_env_var(
+    "SGLANG_MXFP4_USE_OFFICIAL_SHUFFLE", default="true"
+)
+
+
+class PackTopkIds:
+
+    @classmethod
+    def execute(
+        cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor
+    ) -> torch.Tensor:
+        return cls.triton(topk_ids, topk_weights)
+
+    @classmethod
+    def vanilla(
+        cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor
+    ) -> torch.Tensor:
+        weight_bits = (
+            topk_weights.to(torch.bfloat16).view(torch.int16).to(torch.int32) & 0xFFFF
+        )
+        return (topk_ids.to(torch.int32) << 16) | weight_bits
+
+    @classmethod
+    def triton(cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor) -> torch.Tensor:
+        assert (
+            topk_ids.shape == topk_weights.shape
+        ), f"shape mismatch: {topk_ids.shape=} vs {topk_weights.shape=}"
+        assert topk_ids.ndim >= 1, f"expected >=1D, got {topk_ids.shape=}"
+
+        assert (
+            topk_ids.dtype == torch.int32
+        ), f"topk_ids must be int32, got {topk_ids.dtype}"
+        assert (
+            topk_weights.dtype == torch.float32
+        ), f"topk_weights must be float32, got {topk_weights.dtype}"
+
+        assert topk_ids.is_contiguous(), "topk_ids must be contiguous"
+        assert topk_weights.is_contiguous(), "topk_weights must be contiguous"
+
+        out = torch.empty_like(topk_ids, dtype=torch.int32)
+        numel = out.numel()
+        if numel == 0:
+            return out
+
+        BLOCK_SIZE = 1024
+        grid = (triton.cdiv(numel, BLOCK_SIZE),)
+        _pack_topk_ids_triton_kernel[grid](
+            topk_ids,
+            topk_weights,
+            out,
+            numel,
+            BLOCK_SIZE=BLOCK_SIZE,
+        )
+        return out
+
+
+@triton.jit
+def _pack_topk_ids_triton_kernel(
+    topk_ids_ptr,
+    topk_weights_ptr,
+    out_ptr,
+    numel,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < numel
+
+    ids = tl.load(topk_ids_ptr + offsets, mask=mask, other=0)
+    w = tl.load(topk_weights_ptr + offsets, mask=mask, other=0.0)
+
+    w_bf16 = w.to(tl.bfloat16)
+    w_i16 = w_bf16.to(tl.int16, bitcast=True)
+    w_i32 = w_i16.to(tl.int32) & 0xFFFF
+
+    ids_i32 = ids.to(tl.int32)
+    packed = (ids_i32 << 16) | w_i32
+
+    tl.store(out_ptr + offsets, packed, mask=mask)
+
+
+class Mxfp4FlashinferTrtllmMoEMethod:
+
+    def __init__(self, fp8_method, prefix: str):
+        self._fp8 = fp8_method
+        self.prefix = prefix
+        self.flashinfer_mxfp4_moe_precision = (
+            get_global_server_args().flashinfer_mxfp4_moe_precision
+        )
+
+    def create_moe_runner(self, layer, moe_runner_config):
+        self.moe_runner_config = moe_runner_config
+
+        swiglu_limit = moe_runner_config.swiglu_limit
+        assert (
+            swiglu_limit is not None
+        ), f"swiglu_limit must be non-None for DeepSeek V4 (got {swiglu_limit!r})"
+        self._gemm1_clamp_limit_tensor = (
+            torch.full(
+                (layer.num_local_experts,),
+                swiglu_limit,
+                dtype=torch.float32,
+                device=layer.w13_weight.device,
+            )
+            if swiglu_limit is not None
+            else None
+        )
+
+    def create_weights(
+        self,
+        layer,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype,
+        **extra_weight_attrs,
+    ):
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
+
+        fp4_block_k = 32
+
+        w13_weight = Parameter(
+            torch.empty(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size // 2,
+                dtype=torch.int8,
+            ),
+            requires_grad=False,
+        )
+        w2_weight = Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition // 2,
+                dtype=torch.int8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+        layer.register_parameter("w2_weight", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        w13_weight_scale = Parameter(
+            torch.ones(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size // fp4_block_k,
+                dtype=torch.float32,
+            ),
+            requires_grad=False,
+        )
+        w2_weight_scale = Parameter(
+            torch.ones(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition // fp4_block_k,
+                dtype=torch.float32,
+            ),
+            requires_grad=False,
+        )
+        w13_weight_scale.format_ue8m0 = False
+        w2_weight_scale.format_ue8m0 = False
+        scale_attrs = dict(extra_weight_attrs)
+        scale_attrs["quant_method"] = FusedMoeWeightScaleSupported.BLOCK.value
+        layer.register_parameter("w13_weight_scale_inv", w13_weight_scale)
+        set_weight_attrs(w13_weight_scale, scale_attrs)
+        layer.register_parameter("w2_weight_scale_inv", w2_weight_scale)
+        set_weight_attrs(w2_weight_scale, scale_attrs)
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        from sglang.srt.layers.quantization.utils import reorder_w1w3_to_w3w1
+
+        self._fp8.process_weights_after_loading(layer)
+
+        if getattr(layer, "_mega_moe_weights_built", False):
+            return
+
+        w13_w, w13_s = reorder_w1w3_to_w3w1(
+            layer.w13_weight.data, layer.w13_weight_scale_inv.data
+        )
+        layer.w13_weight = Parameter(w13_w, requires_grad=False)
+        layer.w13_weight_scale_inv = Parameter(w13_s, requires_grad=False)
+
+        log_info_on_rank0(
+            logger,
+            f"Shuffling FP4 expert weights for TRT-LLM MxFP4 kernel "
+            f"(layer: {self.prefix})...",
+        )
+
+        w13 = layer.w13_weight.data
+        w2 = layer.w2_weight.data
+        w13_scale = layer.w13_weight_scale_inv.data
+        w2_scale = layer.w2_weight_scale_inv.data
+        num_experts = w13.shape[0]
+
+        if w13_scale.dtype == torch.float32:
+            w13_scale = w13_scale.to(torch.float8_e8m0fnu)
+            w2_scale = w2_scale.to(torch.float8_e8m0fnu)
+
+        epilogue_tile_m = 128
+        g1_w, g1_s, g2_w, g2_s = [], [], [], []
+        if _USE_OFFICIAL_SHUFFLE:
+            cache: dict = {}
+            for i in range(num_experts):
+                w13_u8 = w13[i].view(torch.uint8)
+                w13_s_u8 = w13_scale[i].view(torch.uint8)
+                w2_u8 = w2[i].view(torch.uint8)
+                w2_s_u8 = w2_scale[i].view(torch.uint8)
+
+                perm = _maybe_get_cached_w3_w1_permute_indices(
+                    cache,
+                    w13_u8,
+                    epilogue_tile_m,
+                )
+                g1_w.append(w13_u8[perm.to(w13_u8.device)].contiguous())
+                perm_sf = _maybe_get_cached_w3_w1_permute_indices(
+                    cache,
+                    w13_s_u8,
+                    epilogue_tile_m,
+                    num_elts_per_sf=16,
+                )
+                g1_s.append(
+                    block_scale_interleave(
+                        w13_s_u8[perm_sf.to(w13_s_u8.device)].contiguous()
+                    )
+                )
+
+                perm = get_w2_permute_indices_with_cache(
+                    cache,
+                    w2_u8,
+                    epilogue_tile_m,
+                )
+                g2_w.append(w2_u8[perm.to(w2_u8.device)].contiguous())
+                perm_sf = get_w2_permute_indices_with_cache(
+                    cache,
+                    w2_s_u8,
+                    epilogue_tile_m,
+                    num_elts_per_sf=16,
+                )
+                g2_s.append(
+                    block_scale_interleave(
+                        w2_s_u8[perm_sf.to(w2_s_u8.device)].contiguous()
+                    )
+                )
+        else:
+            for i in range(num_experts):
+                g1_w.append(shuffle_matrix_a(w13[i].view(torch.uint8), epilogue_tile_m))
+                g1_s.append(
+                    shuffle_matrix_sf_a(w13_scale[i].view(torch.uint8), epilogue_tile_m)
+                )
+                g2_w.append(shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m))
+                g2_s.append(
+                    shuffle_matrix_sf_a(w2_scale[i].view(torch.uint8), epilogue_tile_m)
+                )
+
+        layer.w13_weight = Parameter(torch.stack(g1_w), requires_grad=False)
+        layer.w13_weight_scale_inv = Parameter(
+            torch.stack(g1_s)
+            .view(torch.float8_e4m3fn)
+            .reshape(num_experts, w13.shape[1], -1),
+            requires_grad=False,
+        )
+        layer.w2_weight = Parameter(torch.stack(g2_w), requires_grad=False)
+        layer.w2_weight_scale_inv = Parameter(
+            torch.stack(g2_s)
+            .view(torch.float8_e4m3fn)
+            .reshape(num_experts, w2.shape[1], -1),
+            requires_grad=False,
+        )
+
+        self._register_static_scale_ones(layer)
+        torch.cuda.empty_cache()
+
+    def _register_static_scale_ones(self, layer: Module) -> None:
+        device = layer.w13_weight.device
+        for name in (
+            "output1_scale_scalar",
+            "output1_scale_gate_scalar",
+            "output2_scale_scalar",
+        ):
+            layer.register_buffer(
+                name,
+                torch.ones(layer.num_local_experts, device=device, dtype=torch.float32),
+                persistent=False,
+            )
+
+    def apply(
+        self,
+        layer: Module,
+        dispatch_output: DispatchOutput,
+    ) -> CombineInput:
+        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+        from sglang.srt.layers.moe.topk import TopKOutputChecker
+
+        hidden_states = dispatch_output.hidden_states
+        topk_output = dispatch_output.topk_output
+
+        w13 = layer.w13_weight
+        w2 = layer.w2_weight
+        w13_scale = layer.w13_weight_scale_inv
+        w2_scale = layer.w2_weight_scale_inv
+
+        intermediate_size = w2.shape[2] * 2 if w2.dtype == torch.uint8 else w2.shape[2]
+        hidden_size = w13.shape[2] * 2 if w13.dtype == torch.uint8 else w13.shape[2]
+
+        num_local_experts = layer.num_local_experts
+        if w13_scale.dim() == 2:
+            w13_scale = w13_scale.reshape(num_local_experts, 2 * intermediate_size, -1)
+        if w2_scale.dim() == 2:
+            w2_scale = w2_scale.reshape(num_local_experts, hidden_size, -1)
+
+        if TopKOutputChecker.format_is_standard(topk_output):
+            topk_ids = topk_output.topk_ids
+            topk_weights = topk_output.topk_weights
+        elif TopKOutputChecker.format_is_bypassed(topk_output):
+            raise NotImplementedError(
+                "the old code in this branch is WRONG. e.g. it does not consider HashTopK, and may miss args"
+            )
+        else:
+            raise ValueError(f"Unsupported topk output format: {topk_output.format}")
+
+        packed_topk = PackTopkIds.execute(topk_ids, topk_weights)
+
+        precision = self.flashinfer_mxfp4_moe_precision
+        if precision == "bf16":
+            assert hidden_states.dtype == torch.bfloat16
+            x_quant = hidden_states
+            x_scale = None
+            origin_dim = x_quant.shape[-1]
+            if hidden_size != origin_dim:
+                x_quant = torch.nn.functional.pad(
+                    x_quant,
+                    (0, hidden_size - origin_dim),
+                    mode="constant",
+                    value=0.0,
+                )
+        elif precision == "default":
+            x_quant, x_scale = mxfp8_quantize(
+                hidden_states,
+                False,
+                alignment=hidden_size,
+                backend=_MXFP8_QUANTIZE_BACKEND,
+            )
+            x_scale = x_scale.view(torch.float8_e4m3fn).reshape(
+                *hidden_states.shape[:-1], -1
+            )
+        else:
+            raise NotImplementedError(f"Unsupported mxfp4 moe precision: {precision}")
+
+        with use_symmetric_memory(
+            get_tp_group(), disabled=not is_allocation_symmetric()
+        ):
+            num_tokens = x_quant.shape[0]
+            out_hidden_size = (
+                x_quant.shape[-1] * 2
+                if x_quant.dtype == torch.uint8
+                else x_quant.shape[-1]
+            )
+            symm_output = torch.empty(
+                num_tokens, out_hidden_size, dtype=torch.bfloat16, device=x_quant.device
+            )
+
+        output = trtllm_fp4_block_scale_routed_moe(
+            topk_ids=packed_topk,
+            routing_bias=None,
+            hidden_states=x_quant,
+            hidden_states_scale=x_scale,
+            gemm1_weights=w13,
+            gemm1_weights_scale=w13_scale,
+            gemm1_bias=None,
+            gemm1_alpha=None,
+            gemm1_beta=None,
+            gemm1_clamp_limit=self._gemm1_clamp_limit_tensor,
+            gemm2_weights=w2,
+            gemm2_weights_scale=w2_scale,
+            gemm2_bias=None,
+            output1_scale_scalar=layer.output1_scale_scalar,
+            output1_scale_gate_scalar=layer.output1_scale_gate_scalar,
+            output2_scale_scalar=layer.output2_scale_scalar,
+            num_experts=layer.num_experts,
+            top_k=packed_topk.shape[1],
+            n_group=1,
+            topk_group=1,
+            intermediate_size=intermediate_size,
+            local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
+            local_num_experts=num_local_experts,
+            routed_scaling_factor=1.0,
+            routing_method_type=int(RoutingMethodType.TopK),
+            do_finalize=True,
+            tune_max_num_tokens=next_power_of_2(x_quant.shape[0]),
+            output=symm_output,
+        )[0]
+
+        return StandardCombineInput(hidden_states=output)
+
+
+def maybe_fuse_routed_scale_and_shared_add(
+    experts,
+    routed: torch.Tensor,
+    shared: torch.Tensor | None,
+    routed_scaling_factor: float,
+) -> torch.Tensor:
+    # When MxFP4 fusion is on, the upstream `routed *= scale` is skipped and
+    # the scaling is folded into the shared-add via `shared.add_(routed,
+    # alpha=scale)`. With no shared output, the missing scale is applied
+    # in-place. Otherwise `routed` is already scale-final and we just add
+    # `shared` (or pass through if there is none).
+    from sglang.srt.layers.quantization.mxfp4_flashinfer_cutlass_moe import (
+        Mxfp4FlashinferCutlassMoEMethod,
+    )
+    from sglang.srt.layers.quantization.mxfp4_marlin_moe import (
+        Mxfp4MarlinMoEMethod,
+    )
+
+    fused = isinstance(
+        experts.quant_method,
+        (
+            Mxfp4FlashinferTrtllmMoEMethod,
+            Mxfp4FlashinferCutlassMoEMethod,
+            Mxfp4MarlinMoEMethod,
+        ),
+    )
+    if fused:
+        if shared is not None:
+            return shared.add_(routed, alpha=routed_scaling_factor)
+        return routed.mul_(routed_scaling_factor)
+    if shared is not None:
+        routed += shared
+    return routed

python/sglang/srt/layers/quantization/nvfp4_online.py

@@ -0,0 +1,583 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import logging
+import re
+import threading
+from typing import Any, Dict, List, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.quantization.fp8_utils import (
+    block_quant_dequant,
+    inverse_transform_scale_ue8m0,
+)
+from sglang.srt.layers.quantization.modelopt_quant import (
+    ModelOptNvFp4FusedMoEMethod,
+    ModelOptQuantConfig,
+)
+from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+from sglang.srt.layers.quantization.utils import (
+    is_layer_skipped,
+    per_tensor_dequantize,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class NvFp4OnlineConfig(ModelOptQuantConfig):
+    """Config for `--quantization nvfp4_online`.
+
+    This mode is a load-time conversion path, not a serialized NVFP4 checkpoint
+    format. It reuses the ModelOpt NVFP4 MoE parameter layout and fills those
+    parameters by converting BF16/FP16/FP8 expert tensors as they are loaded.
+    Dense layers stay in the source checkpoint precision or quantization path.
+    """
+
+    # Marker consumed by the ModelOpt FP4 layout and the model loader. Serialized
+    # NVFP4 checkpoints use ModelOptFp4Config instead.
+    is_nvfp4_online = True
+    is_checkpoint_nvfp4_serialized = False
+    group_size = 16
+
+    @staticmethod
+    def _normalize_ignored_layers(
+        ignored_layers: Optional[List[str]],
+    ) -> List[str]:
+        if not ignored_layers:
+            return []
+        normalized_ignored_layers = []
+        for layer in ignored_layers:
+            base = layer.removeprefix("model.")
+            normalized_ignored_layers.append(base)
+            normalized_ignored_layers.append(f"model.{base}")
+        return list(dict.fromkeys(normalized_ignored_layers))
+
+    def __init__(
+        self,
+        exclude_modules: Optional[List[str]] = None,
+        packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
+        is_checkpoint_fp8_serialized: bool = False,
+        activation_scheme: str = "dynamic",
+        weight_block_size: Optional[List[int]] = None,
+        use_mxfp8: bool = False,
+    ) -> None:
+        source_ignored_layers = self._normalize_ignored_layers(exclude_modules)
+        fp4_ignored_layers = list(source_ignored_layers)
+        if ignored_layers_str := envs.SGLANG_FP4_IGNORED_LAYERS.get():
+            fp4_ignored_layers.extend(
+                layer.strip()
+                for layer in ignored_layers_str.split(",")
+                if layer.strip()
+            )
+        fp4_ignored_layers = self._normalize_ignored_layers(fp4_ignored_layers)
+        super().__init__(
+            kv_cache_quant_algo=None,
+            exclude_modules=source_ignored_layers,
+            packed_modules_mapping=packed_modules_mapping or {},
+        )
+        self.fp4_ignored_layers = fp4_ignored_layers
+        # Weights use static NVFP4 scales, while FlashInfer computes activation
+        # FP32 scales dynamically per token at runtime.
+        self.use_per_token_activation = True
+        self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
+        self.is_fp4_experts = False
+        self.activation_scheme = activation_scheme
+        self.weight_block_size = weight_block_size
+        self.use_mxfp8 = use_mxfp8
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "nvfp4_online"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.bfloat16, torch.half]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 100
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        return []
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> NvFp4OnlineConfig:
+        quant_method = str(config.get("quant_method", "")).lower()
+        use_mxfp8 = "mxfp8" in quant_method
+        is_checkpoint_fp8_serialized = "fp8" in quant_method or use_mxfp8
+        ignored_layers = config.get("ignored_layers") or config.get(
+            "modules_to_not_convert"
+        )
+        if isinstance(ignored_layers, str):
+            ignored_layers = [ignored_layers]
+        return cls(
+            exclude_modules=ignored_layers,
+            packed_modules_mapping=config.get("packed_modules_mapping"),
+            is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
+            activation_scheme=config.get("activation_scheme", "dynamic"),
+            weight_block_size=config.get("weight_block_size"),
+            use_mxfp8=use_mxfp8,
+        )
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
+        from sglang.srt.layers.linear import LinearBase
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+        from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod, Fp8MoEMethod
+
+        if isinstance(layer, LinearBase):
+            if is_layer_skipped(
+                prefix, self.exclude_modules, self.packed_modules_mapping
+            ) or self.is_layer_excluded(prefix):
+                return UnquantizedLinearMethod()
+            if self.is_checkpoint_fp8_serialized:
+                return Fp8LinearMethod(self)
+            return UnquantizedLinearMethod()
+        if isinstance(layer, FusedMoE):
+            if is_layer_skipped(
+                prefix, self.exclude_modules, self.packed_modules_mapping
+            ) or self.is_layer_excluded(prefix):
+                return None
+            if is_layer_skipped(
+                prefix, self.fp4_ignored_layers, self.packed_modules_mapping
+            ):
+                if self.is_checkpoint_fp8_serialized:
+                    return Fp8MoEMethod(self)
+                return None
+            return ModelOptNvFp4OnlineFusedMoEMethod(self, prefix)
+        return None
+
+
+class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod):
+    """MoE method that converts source expert weights to NVFP4 during loading."""
+
+    def __init__(self, quant_config: NvFp4OnlineConfig, layer_prefix: str):
+        super().__init__(quant_config)
+        self.layer_prefix = layer_prefix
+        layer_match = re.search(r"(?:^|\.)layers\.(\d+)(?:\.|$)", layer_prefix)
+        self.layer_log_name = (
+            f"layer {layer_match.group(1)} ({layer_prefix})"
+            if layer_match is not None
+            else layer_prefix
+        )
+        if not self.enable_flashinfer_trtllm_moe:
+            raise ValueError(
+                "--quantization nvfp4_online supports only "
+                "--moe-runner-backend flashinfer_trtllm or "
+                "flashinfer_trtllm_routed."
+            )
+
+    @staticmethod
+    def _quantize_weight_nvfp4(
+        weight: torch.Tensor,
+        weight_scale_2: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Return packed NVFP4 weight, block scales, and per-tensor decode scale.
+
+        The weight scale is static and per tensor. Callers pass an existing
+        scale when multiple shards must share one global scale, for example the
+        gated w1/w3 pair.
+        """
+        from flashinfer import SfLayout, nvfp4_quantize
+
+        if weight.ndim != 2:
+            raise ValueError(
+                "--quantization nvfp4_online expects 2D expert weights, "
+                f"got shape {tuple(weight.shape)}."
+            )
+        if weight.shape[-1] % 16 != 0:
+            raise ValueError(
+                "--quantization nvfp4_online requires expert weight K to be "
+                f"a multiple of 16, got shape {tuple(weight.shape)}."
+            )
+
+        if weight_scale_2 is None:
+            # weight_scale_2 is the NVFP4 decode scale. FlashInfer consumes its
+            # reciprocal as the global encode scale, matching 448 * 6 / amax.
+            weight_amax = (
+                weight.abs()
+                .nan_to_num()
+                .amax()
+                .to(device=weight.device, dtype=torch.float32)
+            )
+            e4m3_max = (
+                256.0
+                if envs.FLASHINFER_NVFP4_4OVER6.get()
+                and envs.FLASHINFER_NVFP4_4OVER6_E4M3_USE_256.get()
+                else float(torch.finfo(torch.float8_e4m3fn).max)
+            )
+            fp8_fp4_max = e4m3_max * 6.0
+            weight_scale_2 = torch.where(
+                weight_amax > 0,
+                weight_amax / fp8_fp4_max,
+                torch.ones_like(weight_amax),
+            )
+        else:
+            weight_scale_2 = weight_scale_2.to(
+                device=weight.device, dtype=torch.float32
+            )
+        fp4_weight, weight_sf = nvfp4_quantize(
+            weight.contiguous(),
+            1.0 / weight_scale_2,
+            sfLayout=SfLayout.layout_linear,
+            backend="cuda",
+        )
+        rows, cols = weight.shape
+        weight_sf = weight_sf.view(torch.float8_e4m3fn).reshape(rows, cols // 16)
+        return (
+            fp4_weight.reshape(rows, cols // 2),
+            weight_sf.contiguous(),
+            weight_scale_2,
+        )
+
+    @staticmethod
+    def _is_fp8_weight(weight: torch.Tensor) -> bool:
+        fp8_dtypes = {
+            dtype
+            for dtype in (
+                getattr(torch, "float8_e4m3fn", None),
+                getattr(torch, "float8_e5m2", None),
+            )
+            if dtype is not None
+        }
+        return weight.dtype in fp8_dtypes
+
+    @staticmethod
+    def _is_fp8_weight_scale_name(weight_name: str) -> bool:
+        return "weight_scale" in weight_name and "weight_scale_2" not in weight_name
+
+    def _dequantize_fp8_weight(
+        self,
+        weight: torch.Tensor,
+        weight_scale: torch.Tensor,
+        device: torch.device,
+    ) -> torch.Tensor:
+        if self.quant_config.use_mxfp8:
+            raise ValueError(
+                "--quantization nvfp4_online does not support online "
+                "requantization from MXFP8 expert checkpoints."
+            )
+
+        weight = weight.to(device).contiguous()
+        weight_scale = weight_scale.to(device=device).contiguous()
+        if weight_scale.dtype == torch.int32:
+            weight_scale = inverse_transform_scale_ue8m0(
+                weight_scale, mn=weight.shape[-2]
+            )
+        weight_scale = weight_scale.to(dtype=torch.float32).contiguous()
+
+        if weight_scale.numel() == 1 or self.quant_config.weight_block_size is None:
+            return (
+                per_tensor_dequantize(weight, weight_scale)
+                .to(torch.bfloat16)
+                .contiguous()
+            )
+
+        return block_quant_dequant(
+            weight,
+            weight_scale,
+            self.quant_config.weight_block_size,
+            torch.bfloat16,
+        ).contiguous()
+
+    @staticmethod
+    def _should_skip_loaded_expert(
+        layer: torch.nn.Module,
+        param: torch.nn.Parameter,
+        expert_id: Optional[int],
+    ) -> bool:
+        if expert_id is None:
+            return False
+        if getattr(param, "_sglang_require_global_experts", False):
+            return False
+        # With EPLB or explicit expert placement, logical expert IDs can map to
+        # one or more physical experts. Let the canonical MoE loader do that
+        # mapping instead of pre-skipping from the trivial EP layout.
+        from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
+
+        if get_global_expert_location_metadata() is not None:
+            return False
+        return layer._map_global_expert_id_to_local_expert_id(expert_id) == -1
+
+    @staticmethod
+    def _scale_weight_name(weight_name: str) -> str:
+        if "weight" in weight_name:
+            prefix, suffix = weight_name.rsplit("weight", 1)
+            return f"{prefix}weight_scale{suffix}"
+        return f"{weight_name}.weight_scale"
+
+    @staticmethod
+    def _scale_2_weight_name(weight_name: str) -> str:
+        if "weight" in weight_name:
+            prefix, suffix = weight_name.rsplit("weight", 1)
+            return f"{prefix}weight_scale_2{suffix}"
+        return f"{weight_name}.weight_scale_2"
+
+    def get_online_weight_loader(self, layer, original_weight_loader):
+        """Wrap the normal MoE loader with load-time NVFP4 conversion.
+
+        The wrapper quantizes each eligible expert shard as soon as the loader
+        sees enough source data, which avoids materializing and then converting
+        the full checkpoint. FP8 checkpoints stream weight and scale tensors
+        separately, so those pairs are staged until both sides have arrived.
+        """
+        pending_w13_weights = {}
+        pending_w13_lock = threading.Lock()
+        pending_fp8_weights = {}
+        pending_fp8_weight_scales = {}
+        pending_fp8_lock = threading.Lock()
+        quantization_log_lock = threading.Lock()
+        did_log_quantization = False
+
+        def log_quantization_start() -> None:
+            nonlocal did_log_quantization
+            if did_log_quantization:
+                return
+            with quantization_log_lock:
+                if did_log_quantization:
+                    return
+                logger.info(
+                    "Running online NVFP4 quantization for MoE expert weights in %s.",
+                    self.layer_log_name,
+                )
+                did_log_quantization = True
+
+        def store_quantized_weight(
+            param: torch.nn.Parameter,
+            fp4_weight: torch.Tensor,
+            weight_scale: torch.Tensor,
+            weight_scale_2: torch.Tensor,
+            weight_name: str,
+            shard_id: str,
+            expert_id: Optional[int],
+        ) -> None:
+            original_weight_loader(
+                param,
+                fp4_weight,
+                weight_name=weight_name,
+                shard_id=shard_id,
+                expert_id=expert_id,
+            )
+
+            scale_param = (
+                layer.w13_weight_scale
+                if shard_id in ("w1", "w3")
+                else layer.w2_weight_scale
+            )
+            original_weight_loader(
+                scale_param,
+                weight_scale,
+                weight_name=self._scale_weight_name(weight_name),
+                shard_id=shard_id,
+                expert_id=expert_id,
+            )
+            scale_2_param = (
+                layer.w13_weight_scale_2
+                if shard_id in ("w1", "w3")
+                else layer.w2_weight_scale_2
+            )
+            original_weight_loader(
+                scale_2_param,
+                weight_scale_2,
+                weight_name=self._scale_2_weight_name(weight_name),
+                shard_id=shard_id,
+                expert_id=expert_id,
+            )
+
+        def process_loaded_weight(
+            param: torch.nn.Parameter,
+            loaded_weight: torch.Tensor,
+            weight_name: str,
+            shard_id: str,
+            expert_id: Optional[int],
+        ) -> None:
+            log_quantization_start()
+            if shard_id == "w2":
+                loaded_weight = loaded_weight.to(param.device)
+                fp4_weight, weight_scale, weight_scale_2 = self._quantize_weight_nvfp4(
+                    loaded_weight
+                )
+                store_quantized_weight(
+                    param,
+                    fp4_weight,
+                    weight_scale,
+                    weight_scale_2,
+                    weight_name,
+                    shard_id,
+                    expert_id,
+                )
+                return
+
+            pending_key = expert_id
+            current = (
+                param,
+                loaded_weight,
+                weight_name,
+                shard_id,
+                expert_id,
+            )
+            with pending_w13_lock:
+                pending = pending_w13_weights.pop(pending_key, None)
+                if pending is None:
+                    pending_w13_weights[pending_key] = current
+                    return
+
+            (
+                pending_param,
+                pending_weight,
+                pending_name,
+                pending_shard_id,
+                pending_eid,
+            ) = pending
+            if pending_shard_id == shard_id:
+                raise ValueError(
+                    "--quantization nvfp4_online expects paired w1/w3 expert "
+                    f"weights, got two {shard_id} tensors for expert {expert_id}."
+                )
+            pending_weight = pending_weight.to(param.device)
+            loaded_weight = loaded_weight.to(param.device)
+            pending_rows = pending_weight.shape[0]
+            loaded_rows = loaded_weight.shape[0]
+            # Quantize the gated pair together so w1/w3 share one amax-derived
+            # per-tensor FP32 scale, matching the serialized NVFP4 convention.
+            fp4_weight, weight_scale, weight_scale_2 = self._quantize_weight_nvfp4(
+                torch.cat([pending_weight, loaded_weight], dim=0)
+            )
+            pending_fp4_weight, loaded_fp4_weight = fp4_weight.split(
+                [pending_rows, loaded_rows], dim=0
+            )
+            pending_weight_scale, loaded_weight_scale = weight_scale.split(
+                [pending_rows, loaded_rows], dim=0
+            )
+            store_quantized_weight(
+                pending_param,
+                pending_fp4_weight.contiguous(),
+                pending_weight_scale.contiguous(),
+                weight_scale_2,
+                pending_name,
+                pending_shard_id,
+                pending_eid,
+            )
+            store_quantized_weight(
+                param,
+                loaded_fp4_weight.contiguous(),
+                loaded_weight_scale.contiguous(),
+                weight_scale_2,
+                weight_name,
+                shard_id,
+                expert_id,
+            )
+
+        def process_fp8_weight(
+            param: torch.nn.Parameter,
+            loaded_weight: torch.Tensor,
+            weight_name: str,
+            shard_id: str,
+            expert_id: Optional[int],
+        ) -> None:
+            if not self._is_fp8_weight(loaded_weight):
+                process_loaded_weight(
+                    param, loaded_weight, weight_name, shard_id, expert_id
+                )
+                return
+            if not self.quant_config.is_checkpoint_fp8_serialized:
+                raise ValueError(
+                    "--quantization nvfp4_online received an FP8 expert "
+                    "weight, but the checkpoint quantization config does not "
+                    "declare serialized FP8 weights."
+                )
+
+            key = (expert_id, shard_id)
+            with pending_fp8_lock:
+                weight_scale = pending_fp8_weight_scales.pop(key, None)
+                if weight_scale is None:
+                    pending_fp8_weights[key] = (
+                        param,
+                        loaded_weight,
+                        weight_name,
+                        shard_id,
+                        expert_id,
+                    )
+                    return
+
+            log_quantization_start()
+            loaded_weight = self._dequantize_fp8_weight(
+                loaded_weight, weight_scale, param.device
+            )
+            process_loaded_weight(
+                param, loaded_weight, weight_name, shard_id, expert_id
+            )
+
+        def process_fp8_weight_scale(
+            loaded_weight: torch.Tensor,
+            shard_id: str,
+            expert_id: Optional[int],
+        ) -> None:
+            key = (expert_id, shard_id)
+            with pending_fp8_lock:
+                pending = pending_fp8_weights.pop(key, None)
+                if pending is None:
+                    pending_fp8_weight_scales[key] = loaded_weight
+                    return
+
+            log_quantization_start()
+            (
+                pending_param,
+                pending_weight,
+                pending_name,
+                pending_shard_id,
+                pending_eid,
+            ) = pending
+            loaded_weight = self._dequantize_fp8_weight(
+                pending_weight, loaded_weight, pending_param.device
+            )
+            process_loaded_weight(
+                pending_param,
+                loaded_weight,
+                pending_name,
+                pending_shard_id,
+                pending_eid,
+            )
+
+        def nvfp4_online_weight_loader(
+            param: torch.nn.Parameter,
+            loaded_weight: torch.Tensor,
+            weight_name: str,
+            shard_id: str,
+            expert_id: Optional[int],
+        ):
+            if shard_id not in ("w1", "w2", "w3"):
+                original_weight_loader(
+                    param,
+                    loaded_weight,
+                    weight_name=weight_name,
+                    shard_id=shard_id,
+                    expert_id=expert_id,
+                )
+                return
+            if self._should_skip_loaded_expert(layer, param, expert_id):
+                return
+
+            if self._is_fp8_weight_scale_name(weight_name):
+                process_fp8_weight_scale(loaded_weight, shard_id, expert_id)
+                return
+
+            if "weight" in weight_name:
+                process_fp8_weight(
+                    param, loaded_weight, weight_name, shard_id, expert_id
+                )
+                return
+
+            original_weight_loader(
+                param,
+                loaded_weight,
+                weight_name=weight_name,
+                shard_id=shard_id,
+                expert_id=expert_id,
+            )
+
+        return nvfp4_online_weight_loader

python/sglang/srt/layers/utils/common.py

@@ -54,6 +54,41 @@ def copy_or_rebind_param(
         setattr(module, name, Parameter(new_value, requires_grad=False))
 
 
+def alias_or_bind_derived_param(
+    module: torch.nn.Module,
+    source_name: str,
+    derived_name: str,
+    derived_value: torch.Tensor,
+) -> None:
+    """Bind a post-processed (derived) tensor to a derived attribute name.
+
+    When `derived_value` is broadcastable to the source Parameter's shape (and
+    dtype matches), write it broadcast-filled into the source's storage in
+    place and register `derived_name` as an alias of the source Parameter. The
+    two attribute names then share one underlying buffer, so:
+      - apply() can read via `derived_name`
+      - update_weights_from_disk can keep refilling `source_name` (the loader
+        re-runs process_weights_after_loading which re-derives in place)
+      - peak GPU memory is the source size, not source + derived.
+
+    When the shapes are not broadcast-compatible, fall back to allocating a
+    separate Parameter under `derived_name` via copy_or_rebind_param.
+    """
+    derived_value = derived_value.detach()
+    source = getattr(module, source_name, None)
+    if isinstance(source, Parameter) and source.data.dtype == derived_value.dtype:
+        try:
+            broadcast = torch.broadcast_to(derived_value, source.data.shape)
+        except RuntimeError:
+            broadcast = None
+        if broadcast is not None:
+            source.data.copy_(broadcast)
+            source.requires_grad_(False)
+            setattr(module, derived_name, source)
+            return
+    copy_or_rebind_param(module, derived_name, derived_value)
+
+
 class PPMissingLayer(torch.nn.Identity):
     # Adapted from
     # https://github.com/vllm-project/vllm/blob/18ed3132d2bfe1df9a74729457b69243955221e8/vllm/model_executor/models/utils.py#L468C1-L486C1