[Qwen3.8] Enable NVIDIA NVFP4 on DGX Spark with file-backed PLE and PDL router fix (#39126)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: rdxa <rdxa@rdxa-int-spark-01.yvb.moe>
Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: Manrique <nanomlm@gmail.com>
Co-authored-by: yhyang201 <yhyang201@gmail.com>

python/sglang/kernels/jit/csrc/moe/route_radix.cuh

@@ -142,9 +142,10 @@ SGL_DEVICE void route_radix_block(const RouteRadixParams& params, typename Large
     // radix math below is fp32 either way — only the load width differs.
     AlignedVector<packed_t<TScore>, kVecSize / 2> scores_vec;
 
-    // prefetch bias (frozen weight) before the PDL wait
-    bias_vec.load(params.bias, tx);
+    // Bias may be produced by a preceding cast or fill kernel (the caller
+    // does not guarantee a frozen weight), so wait before loading either input.
     PDLWaitPrimary<kUsePDL>();
+    bias_vec.load(params.bias, tx);
     scores_vec.load(scores, tx);
 
 #pragma unroll

python/sglang/kernels/ops/moe/moe_fused_gate.py

@@ -127,17 +127,19 @@ def _router_triton_kernel(
     mask_m = offs_m < M
     mask_n = offs_n < N
 
-    # Prefetch a real bias before the PDL wait. Plain softmax routing has no
-    # bias, so keep the zero value in registers rather than materializing and
-    # clearing a device tensor for every routing call.
+    # PDL may start this grid before prior kernel stores are visible. Bias can
+    # be produced by a preceding cast or fill kernel, so wait before loading
+    # either bias or scores.
+    if USE_PDL:
+        tl.extra.cuda.gdc_wait()
+
+    # Plain softmax routing has no bias, so keep the zero value in registers
+    # rather than materializing and clearing a device tensor per call.
     if HAS_BIAS:
         bias = tl.load(bias_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)
     else:
         bias = tl.zeros([BLOCK_N], dtype=tl.float32)
 
-    if USE_PDL:
-        tl.extra.cuda.gdc_wait()
-
     row_ptr = scores_ptr + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn
     mask2d = mask_m[:, None] & mask_n[None, :]
     scores = tl.load(row_ptr, mask=mask2d, other=0.0).to(

python/sglang/srt/arg_groups/fields/exec_.py

@@ -879,6 +879,30 @@ class ExecOffload(msgspec.Struct):
         ),
     ] = None
 
+    ple_offload_backend: A[
+        str,
+        Arg(
+            help="Host storage for the offloaded Qwen4 PLE n-gram table. "
+            "'pinned' (default) uses CPU pinned memory. 'file' maps a sparse "
+            "file under --ple-offload-dir and lets the gather kernel read it "
+            "directly; use it on unified-memory devices (e.g. GB10 / DGX Spark) "
+            "where pinned host memory comes out of the same pool as the model "
+            "weights. Requires a device that reports "
+            "cudaDevAttrPageableMemoryAccessUsesHostPageTables.",
+            choices=["pinned", "file"],
+        ),
+    ] = "pinned"
+    ple_offload_dir: A[
+        Optional[str],
+        Arg(
+            help="Directory for the file-backed PLE table when "
+            "--ple-offload-backend is 'file'. Defaults to "
+            "$SGLANG_CACHE_DIR/ple/<model path>, one directory per checkpoint. "
+            "The file is sparse and reused across restarts; put it on fast "
+            "local storage (NVMe).",
+        ),
+    ] = None
+
 
 class ExecDllm(msgspec.Struct):
     """Namespace ``exec.dllm``."""

python/sglang/srt/arg_groups/memory_hook.py

@@ -37,6 +37,12 @@ def handle_offload_compatibility(server_args: Any) -> None:
             "would stage the pinned PLE embedding back to the device."
         )
 
+    if cfg.ple_offload_backend == "file" and cfg.ple_offload_embedding is False:
+        raise ValueError(
+            "--ple-offload-backend file requires --ple-offload-embedding: "
+            "the file-backed table is the offloaded table."
+        )
+
 
 def handle_gpu_memory_settings(server_args: Any):
     """

python/sglang/srt/arg_groups/model_overrides/qwen3_moe.py

@@ -10,12 +10,25 @@ from sglang.srt.arg_groups.model_override_base import (
     _register_for,
     resolving_view,
 )
+from sglang.srt.environ import envs
 from sglang.srt.runtime_context import get_platform
 from sglang.srt.utils.common import get_quantization_config
 
 logger = logging.getLogger(__name__)
 
 
+def _mixed_precision_moe_quant_algos(hf_config: Any) -> set:
+    """quant_algo values ModelOpt MIXED_PRECISION assigns to `*.experts` layers."""
+    quantization_config = getattr(hf_config, "quantization_config", None)
+    if not isinstance(quantization_config, dict):
+        return set()
+    return {
+        str(info.get("quant_algo", "")).upper()
+        for name, info in quantization_config.get("quantized_layers", {}).items()
+        if ".experts" in name and isinstance(info, dict)
+    }
+
+
 @_register_for(
     "Qwen3MoeForCausalLM",
     "Qwen3VLMoeForConditionalGeneration",
@@ -38,8 +51,38 @@ def _qwen3_moe_family_overrides(server_args: Any, hf_config: Any) -> dict:
         ):
             overrides["quantization"] = quant_method
             quantization = quant_method
-        if (
-            (quantization in ("fp8", "modelopt_fp4") or quantization is None)
+        has_w4a16_moe_layers = (
+            quantization == "modelopt_mixed"
+            and "W4A16_NVFP4" in _mixed_precision_moe_quant_algos(hf_config)
+        )
+        if has_w4a16_moe_layers:
+            # trtllm-gen only has the W4A4 NVFP4 MoE path.
+            # CuTe DSL v2 also supports W4A16 with BF16 activations when opted in.
+            use_cutedsl_w4a16 = (
+                cfg.moe_runner_backend == "flashinfer_cutedsl"
+                and cfg.moe_a2a_backend in ("none", "flashinfer")
+                and envs.SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16.get()
+            )
+            if (
+                cfg.moe_runner_backend not in ("auto", "marlin")
+                and not use_cutedsl_w4a16
+            ):
+                raise ValueError(
+                    "W4A16_NVFP4 MoE layers require --moe-runner-backend=marlin, "
+                    "or flashinfer_cutedsl with --moe-a2a-backend=none/flashinfer "
+                    "and SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1."
+                )
+            if cfg.moe_runner_backend == "auto":
+                overrides["moe_runner_backend"] = "marlin"
+                logger.info(
+                    "Use marlin as MoE runner backend for "
+                    f"{hf_config.architectures[0]} with W4A16_NVFP4 MoE layers"
+                )
+        elif (
+            (
+                quantization in ("fp8", "modelopt_fp4", "modelopt_mixed")
+                or quantization is None
+            )
             and cfg.moe_a2a_backend == "none"
             and cfg.moe_runner_backend == "auto"
         ):

python/sglang/srt/configs/qwen4_exp.py

@@ -31,6 +31,8 @@ class Qwen4ExpTextConfig(Qwen3NextConfig):
         ngram_vocab_size_base=20000000,
         make_ngram_vocab_size_divisible_by=128,
         ple_offload_embedding=False,
+        ple_offload_backend="pinned",
+        ple_offload_dir=None,
         ple_embedding_dtype=None,
         index_share_for_mtp_iteration=True,
         rope_parameters=None,
@@ -68,6 +70,10 @@ class Qwen4ExpTextConfig(Qwen3NextConfig):
         self.ngram_vocab_size_base = ngram_vocab_size_base
         self.make_ngram_vocab_size_divisible_by = make_ngram_vocab_size_divisible_by
         self.ple_offload_embedding = ple_offload_embedding
+        # Host storage for the offloaded table: "pinned" or "file" (a sparse
+        # file-backed mmap for unified-memory devices); see --ple-offload-backend.
+        self.ple_offload_backend = ple_offload_backend
+        self.ple_offload_dir = ple_offload_dir
         # "float8_e4m3fn" keeps fp8 PLE tables fp8-resident; text_config-scoped.
         self.ple_embedding_dtype = ple_embedding_dtype
         # Draft decode steps reuse the draft-extend indexer top-k (IndexShare).

python/sglang/srt/environ.py

@@ -304,6 +304,18 @@ class Envs:
     # Bitwise-exact, shape-guarded Qwen4 PLE decode fusion. Unsupported inputs
     # and phases fall back to the original implementation.
     SGLANG_ENABLE_QWEN4_PLE_FUSION = EnvBool(True)
+    # --ple-offload-backend file: where the sparse, file-backed PLE table lives
+    # (deterministic name, reused across restarts), whether prefill-sized
+    # gathers hint the page cache first, and an escape hatch for the device
+    # attribute check (pageable host memory reachable through host page tables).
+    SGLANG_QWEN4_PLE_FILE_DIR = EnvStr(lambda: _default_cache_subdir("ple"))
+    SGLANG_QWEN4_PLE_FILE_PREFETCH = EnvBool(True)
+    SGLANG_QWEN4_PLE_FILE_SKIP_DEVICE_CHECK = EnvBool(False)
+    # Faulting rows in maps whole page-cache folios, so the mapping creeps
+    # towards full residency (~45 KB/token) and eats the free memory that
+    # sizes the KV pool. Cap its resident set; 0 disables the trim.
+    SGLANG_QWEN4_PLE_FILE_RSS_BUDGET_GB = EnvFloat(8.0)
+    SGLANG_QWEN4_PLE_FILE_RSS_INTERVAL_S = EnvFloat(30.0)
     SGLANG_PREFETCH_BLOCK_SIZE_MB = EnvInt(16)
     SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION = EnvBool(False)
     SGLANG_ENABLE_WEIGHT_LOADER_V2 = EnvBool(False)

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

@@ -2506,6 +2506,18 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             else:
                 moe_runner_backend = MoeRunnerBackend.TRITON
 
+        if (
+            moe_runner_backend.is_flashinfer_cutlass()
+            or moe_runner_backend.is_flashinfer_cutedsl()
+        ):
+            # Neither runner has an fp8 MoE path; they get pinned globally for
+            # NVFP4 experts on sm120, so run this layer's fp8 experts on triton.
+            logger.info(
+                "Fp8MoEMethod has no %s path; using triton for its fp8 experts.",
+                moe_runner_backend.name,
+            )
+            moe_runner_backend = MoeRunnerBackend.TRITON
+
         if (
             moe_runner_backend.is_deep_gemm()
             or moe_runner_backend.is_triton()

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

@@ -802,12 +802,14 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
         nvfp4_config: ModelOptFp4Config,
         nvfp4a16_config: ModelOptFp4Config,
         mxfp8_config: Fp8Config,
+        fp8_block_config: Fp8Config,
     ) -> None:
         super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping)
         self.quantized_layers = quantized_layers
         self.fp8_config = fp8_config
         self.fp8_pb_wo_config = fp8_pb_wo_config
         self.mxfp8_config = mxfp8_config
+        self.fp8_block_config = fp8_block_config
         self.nvfp4_config = nvfp4_config
         self.nvfp4a16_config = nvfp4a16_config
 
@@ -864,6 +866,9 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
             exclude_modules = quantization_section.get("exclude_modules")
             quantized_layers = quantization_section.get("quantized_layers", {})
 
+        # ModelOpt emits `ignore: []` or omits it; is_layer_skipped iterates it.
+        exclude_modules = list(exclude_modules or [])
+
         if quant_algo != "MIXED_PRECISION":
             raise ValueError(
                 "ModelOptMixedPrecisionConfig only supports MIXED_PRECISION checkpoints."
@@ -904,6 +909,13 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
             packed_modules_mapping=packed_modules_mapping,
             use_mxfp8=True,
         )
+        # ModelOpt FP8_BLOCK_SCALES: 128x128 block fp8 with weight_scale_inv.
+        fp8_block_config = Fp8Config(
+            is_checkpoint_fp8_serialized=True,
+            activation_scheme="dynamic",
+            weight_block_size=[128, 128],
+            packed_modules_mapping=packed_modules_mapping,
+        )
         nvfp4_config = ModelOptFp4Config(
             is_checkpoint_nvfp4_serialized=True,
             kv_cache_quant_algo=kv_cache_quant_algo,
@@ -928,6 +940,7 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
             fp8_config=fp8_config,
             fp8_pb_wo_config=fp8_pb_wo_config,
             mxfp8_config=mxfp8_config,
+            fp8_block_config=fp8_block_config,
             nvfp4_config=nvfp4_config,
             nvfp4a16_config=nvfp4a16_config,
         )
@@ -986,9 +999,17 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
             candidates.append(
                 "language_model.model." + prefix[len("model.language_model.") :]
             )
+            candidates.append("model." + prefix[len("model.language_model.") :])
+        elif prefix.startswith("model."):
+            # VL models such as Qwen4-Exp name the text stack `model.layers.*`
+            # while ModelOpt keys it `model.language_model.layers.*`.
+            candidates.append("model.language_model." + prefix[len("model.") :])
 
         return tuple(dict.fromkeys(candidates))
 
+    def resolve_quant_algo(self, prefix: str) -> Optional[str]:
+        return self._resolve_quant_algo(prefix)
+
     def get_quant_method(
         self, layer: torch.nn.Module, prefix: str
     ) -> Optional[QuantizeMethodBase]:
@@ -1010,6 +1031,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
                 return ModelOptFp8LinearMethod(self.fp8_config)
             if quant_algo == "FP8_PB_WO":
                 return Fp8LinearMethod(self.fp8_pb_wo_config)
+            if quant_algo == "FP8_BLOCK_SCALES":
+                return Fp8LinearMethod(self.fp8_block_config)
             if quant_algo == "MXFP8":
                 return Fp8LinearMethod(self.mxfp8_config)
             if quant_algo == "NVFP4":
@@ -1039,6 +1062,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
                 return ModelOptFp8MoEMethod(self.fp8_config)
             if quant_algo == "MXFP8":
                 return Fp8MoEMethod(self.mxfp8_config)
+            if quant_algo == "FP8_BLOCK_SCALES":
+                return Fp8MoEMethod(self.fp8_block_config)
             if quant_algo == "NVFP4":
                 return ModelOptNvFp4FusedMoEMethod(self.nvfp4_config)
             if quant_algo == "W4A16_NVFP4":

python/sglang/srt/model_executor/model_runner_components/load_model_utils.py

@@ -275,6 +275,27 @@ def load_model_with_memory_saver(
             )
         if is_qwen4_exp:
             model_config.hf_text_config.ple_offload_embedding = ple_offload_embedding
+            model_config.hf_text_config.ple_offload_backend = (
+                get_exec().offload.ple_offload_backend
+            )
+            if get_exec().offload.ple_offload_backend != "file":
+                model_config.hf_text_config.ple_offload_dir = (
+                    get_exec().offload.ple_offload_dir
+                )
+            else:
+                from sglang.srt.models.qwen4_exp_ple_table import (
+                    check_file_backend_supported,
+                    default_ple_table_dir,
+                )
+
+                model_config.hf_text_config.ple_offload_dir = (
+                    get_exec().offload.ple_offload_dir
+                    or default_ple_table_dir(get_model().model_path)
+                )
+                if ple_offload_embedding and device == "cuda":
+                    check_file_backend_supported(
+                        torch.cuda.current_device() if torch.cuda.is_available() else 0
+                    )
 
     enable_cpu_backup = get_exec().features.enable_weights_cpu_backup or (
         is_draft_worker and get_exec().features.enable_draft_weights_cpu_backup

python/sglang/srt/models/qwen3_5_mtp.py

@@ -58,12 +58,14 @@ def _mtp_quant_config(quant_config):
     # Serialized Qwen3.5 ModelOpt checkpoints keep embedded MTP weights in
     # BF16. Disable quantization for those checkpoints; non-serialized
     # modelopt_fp4 still converts MoE expert weights on load.
+    if quant_config and quant_config.get_name() == "modelopt_mixed":
+        # MIXED_PRECISION lists mtp.* layers only when the MTP head is quantized.
+        if any(name.startswith("mtp.") for name in quant_config.quantized_layers):
+            return quant_config
+        return None
     if quant_config and (
-        quant_config.get_name() == "modelopt_mixed"
-        or (
-            quant_config.get_name() == "modelopt_fp4"
-            and quant_config.is_checkpoint_nvfp4_serialized
-        )
+        quant_config.get_name() == "modelopt_fp4"
+        and quant_config.is_checkpoint_nvfp4_serialized
     ):
         return None
     if is_npu() and get_spec().speculative_draft_model_quantization is None:

python/sglang/srt/models/qwen4_exp.py

@@ -43,6 +43,9 @@ from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
+from sglang.srt.layers.quantization.modelopt_quant import (
+    ModelOptMixedPrecisionConfig,
+)
 from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod
 from sglang.srt.layers.utils import get_layer_id
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
@@ -60,6 +63,11 @@ from sglang.srt.models.qwen3_5 import (
     Qwen3_5LinearDecoderLayer,
 )
 from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration
+from sglang.srt.models.qwen4_exp_ple_table import (
+    allocate_ple_host_table,
+    make_ple_file_prefetcher,
+    make_ple_file_rss_trimmer,
+)
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import logger
 
@@ -68,6 +76,24 @@ from sglang.srt.utils import logger
 _QSA_INDEXER_OVERLAP_TOKEN_THRESHOLD = 1024
 
 
+def _ple_table_is_fp8(
+    config: Qwen4ExpTextConfig,
+    quant_config: Optional[QuantizationConfig],
+    prefix: str,
+) -> bool:
+    """fp8 PLE shards: declared by config, an fp8 checkpoint, or a ModelOpt
+    MIXED_PRECISION entry for the ngram table (nvidia/*-Flash-Next-NVFP4)."""
+    if config.ple_embedding_dtype == "float8_e4m3fn":
+        return True
+    if quant_config is None:
+        return False
+    if quant_config.get_name() == "fp8":
+        return True
+    if isinstance(quant_config, ModelOptMixedPrecisionConfig):
+        return quant_config.resolve_quant_algo(prefix) == "FP8"
+    return False
+
+
 def _get_ple_forward_mode(forward_batch: ForwardBatch) -> ForwardMode:
     if forward_batch._original_forward_mode is not None:
         return forward_batch._original_forward_mode
@@ -423,6 +449,7 @@ class Qwen4ExpNGramEmbedding(nn.Module):
         embedding_dim: int,
         ple_layer_index: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.config = config
@@ -479,13 +506,13 @@ class Qwen4ExpNGramEmbedding(nn.Module):
             and get_attention_dp_size() > 1
             and not self.use_attn_tp_ngram
         )
+        ngram_prefix = f"{prefix}.ngram_embedding" if prefix else "ngram_embedding"
         self.ngram_embedding = VocabParallelEmbedding(
             padded_vocab_size,
             self.head_dim_per_ngram,
             params_dtype=(
                 torch.float8_e4m3fn
-                if (quant_config is not None and quant_config.get_name() == "fp8")
-                or getattr(config, "ple_embedding_dtype", None) == "float8_e4m3fn"
+                if _ple_table_is_fp8(config, quant_config, ngram_prefix)
                 else torch.bfloat16
             ),
             output_dtype=torch.bfloat16,
@@ -739,7 +766,7 @@ def _gather_ple_embedding_from_pinned_kernel(
 
 
 class Qwen4ExpPinnedHostEmbedding(VocabParallelEmbedding):
-    """PLE table read directly from pinned host memory.
+    """PLE table read directly from host memory (pinned, or a file-backed mmap).
 
     The table stays in its checkpoint storage dtype (fp8 with a per-tensor
     weight_scale for fp8 checkpoints, bf16 otherwise); gathers emit bf16.
@@ -764,7 +791,13 @@ class Qwen4ExpPinnedHostEmbedding(VocabParallelEmbedding):
         "num_added_embeddings_per_partition",
     )
 
-    def __init__(self, embedding: VocabParallelEmbedding) -> None:
+    def __init__(
+        self,
+        embedding: VocabParallelEmbedding,
+        *,
+        backend: str = "pinned",
+        table_dir: Optional[str] = None,
+    ) -> None:
         nn.Module.__init__(self)
         if not isinstance(embedding.quant_method, UnquantizedEmbeddingMethod):
             raise NotImplementedError(
@@ -786,15 +819,23 @@ class Qwen4ExpPinnedHostEmbedding(VocabParallelEmbedding):
         self.quant_method = None
 
         source_weight = embedding.weight
-        cpu_weight = nn.Parameter(
-            torch.empty(
-                source_weight.shape,
-                dtype=source_weight.dtype,
-                device="cpu",
-                pin_memory=True,
+        host_table = allocate_ple_host_table(
+            shape=source_weight.shape,
+            dtype=source_weight.dtype,
+            backend=backend,
+            table_dir=table_dir,
+            # Each TP rank holds a different vocabulary shard of the same shape.
+            tag=(
+                f"rows{self.shard_indices.org_vocab_start_index}"
+                f"-{self.shard_indices.org_vocab_end_index}"
             ),
-            requires_grad=False,
         )
+        # Only the file backend has anything to prefetch (rows live on storage).
+        self._file_prefetcher = make_ple_file_prefetcher(host_table)
+        # ... and only it needs its resident set bounded: a fault maps a whole
+        # folio, so the mapping would otherwise creep towards the full table.
+        self._file_rss_trimmer = make_ple_file_rss_trimmer(host_table)
+        cpu_weight = nn.Parameter(host_table, requires_grad=False)
         for name, value in vars(source_weight).items():
             setattr(cpu_weight, name, value)
         cpu_weight.weight_loader = self.weight_loader
@@ -837,6 +878,12 @@ class Qwen4ExpPinnedHostEmbedding(VocabParallelEmbedding):
 
         flat_ids = input_ids.reshape(-1).long()
         if flat_ids.numel():
+            if self._file_prefetcher is not None:
+                self._file_prefetcher.enqueue(
+                    flat_ids,
+                    vocab_start=self.shard_indices.org_vocab_start_index,
+                    vocab_end=self.shard_indices.org_vocab_end_index,
+                )
             _gather_ple_embedding_from_pinned_kernel[(flat_ids.numel(),)](
                 self.weight.data_ptr(),
                 flat_ids,
@@ -881,10 +928,13 @@ class Qwen4ExpPLELayer(nn.Module):
             self.ple_embed_dim,
             ple_layer_index=ple_layer_index,
             quant_config=quant_config,
+            prefix=f"{prefix}.ple_embedding" if prefix else "ple_embedding",
         )
         if config.ple_offload_embedding:
             self.ple_embedding.ngram_embedding = Qwen4ExpPinnedHostEmbedding(
-                self.ple_embedding.ngram_embedding
+                self.ple_embedding.ngram_embedding,
+                backend=getattr(config, "ple_offload_backend", "pinned"),
+                table_dir=getattr(config, "ple_offload_dir", None),
             )
         self.short_conv_dilation = self.ple_embedding.ngram_size
         self.short_conv_state_len = (

python/sglang/srt/models/qwen4_exp_ple_table.py

@@ -0,0 +1,483 @@
+"""Host-side storage for the offloaded Qwen4-Exp PLE n-gram table.
+
+``--ple-offload-embedding`` keeps the PLE table (47.7 GiB in fp8 for
+Qwen3.8-Flash-Next) out of device memory and lets the Triton gather kernel read
+rows straight from a host pointer. Two backends provide that pointer:
+
+``pinned`` (default)
+    ``torch.empty(..., pin_memory=True)``. On a discrete GPU this frees VRAM.
+
+``file``
+    A file-backed, shared ``mmap`` of a sparse file under
+    ``--ple-offload-dir``. Meant for unified-memory parts (GB10 / DGX Spark and
+    similar), where pinned host memory comes out of the *same* pool as the
+    model weights and ``pinned`` therefore frees nothing: Qwen3.8-Flash-Next is
+    126.0 GiB of weights on a 121.63 GiB box and does not boot with ``pinned``.
+    The kernel dereferences the pageable pointer directly, which only works on
+    devices that report ``cudaDevAttrPageableMemoryAccessUsesHostPageTables``;
+    rows are paged in from storage on demand, the file is sparse, deterministic
+    in name and reused across restarts, and gathers of prefill size hint the
+    page cache (``posix_fadvise(WILLNEED)``) so page faults are served
+    concurrently instead of one at a time. A background trimmer keeps the
+    mapping's resident set under a budget, because faulting rows in maps whole
+    page-cache folios and the table would otherwise creep towards full
+    residency (see ``PleFileRssTrimmer``).
+
+This module has no Triton or CUDA-kernel imports so that its allocator and
+prefetcher can be unit-tested on CPU.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import ctypes.util
+import logging
+import os
+import re
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from typing import Optional, Sequence
+
+import torch
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+_LIBC: Optional[ctypes.CDLL] = None
+_SMAPS_HEADER = re.compile(r"^([0-9a-f]+)-([0-9a-f]+) ")
+_SMAPS_RSS = re.compile(r"^Rss:\s+(\d+) kB")
+
+PLE_OFFLOAD_BACKENDS = ("pinned", "file")
+
+# cudaDeviceAttr enum values (cuda_runtime_api.h).
+_CUDA_DEV_ATTR_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES = 100
+_MADV_RANDOM = 1
+_MADV_DONTNEED = 4
+_PAGE_SHIFT = 12
+# One MADV_DONTNEED call takes mmap_lock for its whole range; over the full
+# 47.7 GiB table that is ~3.5 s during which every fault in the process --
+# including the ones the gather kernel takes -- stalls. Trim in slices.
+PLE_FILE_RSS_TRIM_CHUNK_BYTES = 1 << 30
+# Below this many rows a gather is decode-sized (16 rows per token): the page
+# faults are cheap and the host-side hint would cost more than it saves.
+PLE_FILE_PREFETCH_MIN_ROWS = 2048
+
+
+class PleFilePrefetcher:
+    """Hint the page cache about the rows a prefill-sized gather is about to read.
+
+    With the table on storage, a cold prefill chunk faults tens of thousands of
+    4 KiB pages one at a time from inside the gather kernel. Advising them
+    first (``posix_fadvise(WILLNEED)`` per distinct page, on one background
+    thread) lets the block layer serve them concurrently. Measured on a GB10 /
+    NVMe: cold prefill 650-750 tok/s -> 1,000-2,100 tok/s (warm: ~2,200-2,600).
+    Decode-sized gathers are skipped; nothing runs during CUDA-graph capture.
+    """
+
+    def __init__(
+        self,
+        path: str,
+        row_bytes: int,
+        min_rows: int = PLE_FILE_PREFETCH_MIN_ROWS,
+    ) -> None:
+        self._fd = os.open(path, os.O_RDONLY)
+        self._row_bytes = int(row_bytes)
+        self._min_rows = int(min_rows)
+        self._pool = ThreadPoolExecutor(max_workers=1)
+
+    @staticmethod
+    def pages_for_rows(row_ids: torch.Tensor, row_bytes: int) -> list[int]:
+        start = row_ids.to(torch.int64) * row_bytes
+        end = start + (row_bytes - 1)
+        return (
+            torch.cat([start >> _PAGE_SHIFT, end >> _PAGE_SHIFT])
+            .unique(sorted=True)
+            .tolist()
+        )
+
+    def _advise(self, pages: list[int]) -> None:
+        for p in pages:
+            try:
+                os.posix_fadvise(
+                    self._fd, p << _PAGE_SHIFT, 1 << _PAGE_SHIFT, os.POSIX_FADV_WILLNEED
+                )
+            except OSError:
+                return
+
+    def enqueue(
+        self,
+        flat_ids: torch.Tensor,
+        *,
+        vocab_start: int = 0,
+        vocab_end: Optional[int] = None,
+    ) -> bool:
+        """Queue the hint for ``flat_ids``. Returns whether anything was queued."""
+        if flat_ids.numel() < self._min_rows:
+            return False
+        if flat_ids.is_cuda and torch.cuda.is_current_stream_capturing():
+            return False
+        # The .cpu() syncs the stream; acceptable for prefill chunks (~1 s) and
+        # it is what lets the page set be computed without touching the kernel.
+        row_ids = flat_ids.detach().cpu()
+        if vocab_end is not None:
+            # The file contains only this rank's vocabulary shard.
+            row_ids = row_ids[(row_ids >= vocab_start) & (row_ids < vocab_end)]
+        row_ids = row_ids - vocab_start
+        if row_ids.numel() == 0:
+            return False
+        pages = self.pages_for_rows(row_ids, self._row_bytes)
+        self._pool.submit(self._advise, pages)
+        return True
+
+    def close(self) -> None:
+        self._pool.shutdown(wait=False)
+        try:
+            os.close(self._fd)
+        except OSError:
+            pass
+
+
+class PleFileRssTrimmer:
+    """Keep the mapped table's resident set under a budget.
+
+    Every random row fault maps in a whole page-cache folio, so with large
+    folios (Linux 6.x) the mapping's Rss climbs towards the table's full size
+    while a generated token only reads a few KB of it: measured ~45 KB of Rss
+    growth per token on a GB10. On a unified-memory part that is not a slow
+    leak, it is a countdown -- the free-memory readings that size the KV pool
+    come from the same pool the folios are accumulating in.
+
+    ``MADV_RANDOM`` does not prevent it (it limits readahead I/O, not the
+    mapping-in of folios already in cache) and ``posix_fadvise(DONTNEED)`` does
+    not release them either. ``MADV_DONTNEED`` over the mapping does: the page
+    table entries go, the pages stay in the page cache, and hot rows come back
+    at minor-fault cost.
+
+    Dropping entries under a running gather is the state this backend already
+    handles: the file starts out entirely unfaulted and every cold row is
+    faulted in from inside the kernel through the same host page tables. What
+    must not happen is one ``madvise`` call over the whole table, so the trim
+    is chunked (see ``PLE_FILE_RSS_TRIM_CHUNK_BYTES``) and runs on its own
+    daemon thread -- decode replays a CUDA graph and executes no Python, so a
+    hook in the gather would never fire in the phase that grows the table.
+    """
+
+    def __init__(
+        self,
+        addr: int,
+        nbytes: int,
+        budget_bytes: int,
+        interval_s: float,
+        chunk_bytes: int = PLE_FILE_RSS_TRIM_CHUNK_BYTES,
+    ) -> None:
+        self._addr = int(addr)
+        self._nbytes = int(nbytes)
+        self._budget = int(budget_bytes)
+        self._interval = float(interval_s)
+        self._chunk = int(chunk_bytes)
+        self._stop = threading.Event()
+        self._thread = threading.Thread(
+            target=self._loop, name="ple-file-rss-trim", daemon=True
+        )
+
+    def start(self) -> None:
+        self._thread.start()
+
+    def mapping_rss_bytes(self) -> Optional[int]:
+        """Resident bytes of the VMAs backing the table, or None off Linux."""
+        return _mapping_rss_bytes(self._addr, self._nbytes)
+
+    def trim_once(self) -> int:
+        """Drop the mapping's resident pages if over budget. Returns bytes freed."""
+        before = self.mapping_rss_bytes()
+        if before is None or before <= self._budget:
+            return 0
+        for offset in range(0, self._nbytes, self._chunk):
+            if self._stop.is_set():
+                break
+            length = min(self._chunk, self._nbytes - offset)
+            if not _madvise(self._addr + offset, length, _MADV_DONTNEED):
+                return 0
+            # Let the faults that queued behind mmap_lock through.
+            self._stop.wait(0.005)
+        after = self.mapping_rss_bytes()
+        freed = before - after if after is not None else 0
+        logger.info(
+            "PLE table: trimmed resident set %.1f -> %.1f GiB (budget %.1f GiB)",
+            before / 2**30,
+            (after if after is not None else 0) / 2**30,
+            self._budget / 2**30,
+        )
+        return max(freed, 0)
+
+    def _loop(self) -> None:
+        while not self._stop.wait(self._interval):
+            try:
+                self.trim_once()
+            except Exception as exc:  # advisory only; never fail a request
+                logger.warning("PLE table: resident-set trim skipped (%s)", exc)
+
+    def close(self) -> None:
+        self._stop.set()
+
+
+def allocate_ple_host_table(
+    shape: Sequence[int],
+    dtype: torch.dtype,
+    backend: str = "pinned",
+    table_dir: Optional[str] = None,
+    tag: Optional[str] = None,
+) -> torch.Tensor:
+    """Return a host tensor of ``shape``/``dtype`` for the PLE table.
+
+    For the file backend, ``table_dir`` should be private to one checkpoint
+    (the server defaults it to ``$SGLANG_CACHE_DIR/ple/<model path>``): the
+    file name only encodes shape, dtype and ``tag``, and every boot rewrites
+    the whole table through the weight loader.
+    """
+    if backend not in PLE_OFFLOAD_BACKENDS:
+        raise ValueError(
+            f"unknown PLE offload backend {backend!r}; choose from {PLE_OFFLOAD_BACKENDS}"
+        )
+    if backend == "pinned":
+        return torch.empty(tuple(shape), dtype=dtype, device="cpu", pin_memory=True)
+
+    numel = 1
+    for d in shape:
+        numel *= int(d)
+    nbytes = numel * torch.empty(0, dtype=dtype).element_size()
+    table_dir = os.path.expanduser(table_dir or envs.SGLANG_QWEN4_PLE_FILE_DIR.get())
+    os.makedirs(table_dir, exist_ok=True)
+    path = os.path.join(table_dir, ple_table_file_name(shape, dtype, tag))
+    if not os.path.exists(path) or os.path.getsize(path) != nbytes:
+        # Sparse: only pages that get written take disk space.
+        with open(path, "wb") as f:
+            f.truncate(nbytes)
+    logger.info(
+        "PLE table: file-backed mmap %s (%.1f GiB, %s)", path, nbytes / 2**30, dtype
+    )
+    storage = torch.from_file(path, shared=True, size=nbytes, dtype=torch.uint8)
+    _madvise_random(storage, nbytes)
+    table = storage.view(dtype).view(*[int(d) for d in shape])
+    table._sglang_ple_file_path = path  # consumed by PleFilePrefetcher
+    return table
+
+
+def make_ple_file_prefetcher(table: torch.Tensor) -> Optional[PleFilePrefetcher]:
+    """A prefetcher for a table returned by ``allocate_ple_host_table(..., "file")``."""
+    path = getattr(table, "_sglang_ple_file_path", None)
+    if path is None or not envs.SGLANG_QWEN4_PLE_FILE_PREFETCH.get():
+        return None
+    row_bytes = (
+        int(table.shape[-1]) * table.element_size()
+        if table.dim() >= 2
+        else table.element_size()
+    )
+    prefetcher = PleFilePrefetcher(path=path, row_bytes=row_bytes)
+    logger.info(
+        "PLE table: WILLNEED prefetch on for gathers of >= %d rows (row = %d B)",
+        PLE_FILE_PREFETCH_MIN_ROWS,
+        row_bytes,
+    )
+    return prefetcher
+
+
+def make_ple_file_rss_trimmer(table: torch.Tensor) -> Optional[PleFileRssTrimmer]:
+    """A started trimmer for a table from ``allocate_ple_host_table(..., "file")``.
+
+    ``SGLANG_QWEN4_PLE_FILE_RSS_BUDGET_GB=0`` turns it off; it is also absent
+    where the resident set cannot be read (no ``/proc/self/smaps``).
+    """
+    path = getattr(table, "_sglang_ple_file_path", None)
+    if path is None:
+        return None
+    budget_gb = float(envs.SGLANG_QWEN4_PLE_FILE_RSS_BUDGET_GB.get())
+    if budget_gb <= 0:
+        return None
+    nbytes = table.numel() * table.element_size()
+    if _mapping_rss_bytes(table.data_ptr(), nbytes) is None:
+        logger.warning(
+            "PLE table: resident-set trim off, /proc/self/smaps is not readable; "
+            "the mapping will creep towards %.1f GiB resident",
+            nbytes / 2**30,
+        )
+        return None
+    trimmer = PleFileRssTrimmer(
+        addr=table.data_ptr(),
+        nbytes=nbytes,
+        budget_bytes=int(budget_gb * 2**30),
+        interval_s=float(envs.SGLANG_QWEN4_PLE_FILE_RSS_INTERVAL_S.get()),
+    )
+    trimmer.start()
+    logger.info(
+        "PLE table: resident set capped at %.1f GiB, checked every %.0f s",
+        budget_gb,
+        float(envs.SGLANG_QWEN4_PLE_FILE_RSS_INTERVAL_S.get()),
+    )
+    return trimmer
+
+
+def check_file_backend_supported(device_index: int = 0) -> None:
+    """Fail fast at load time instead of silently reading garbage in the kernel."""
+    if envs.SGLANG_QWEN4_PLE_FILE_SKIP_DEVICE_CHECK.get():
+        logger.warning(
+            "PLE table: file backend device check skipped by "
+            "SGLANG_QWEN4_PLE_FILE_SKIP_DEVICE_CHECK"
+        )
+        return
+    supported = device_uses_host_page_tables(device_index)
+    if supported is None:
+        raise RuntimeError(
+            "--ple-offload-backend file: could not query "
+            "cudaDevAttrPageableMemoryAccessUsesHostPageTables. Set "
+            "SGLANG_QWEN4_PLE_FILE_SKIP_DEVICE_CHECK=1 only if you know the "
+            "device reads pageable host memory through the host page tables."
+        )
+    if not supported:
+        raise ValueError(
+            "--ple-offload-backend file needs a device whose pageable host "
+            "memory accesses go through the host page tables (unified-memory "
+            "parts such as GB10). This device reports it does not; use "
+            "--ple-offload-backend pinned."
+        )
+
+
+def default_ple_table_dir(model_path: str) -> str:
+    """``$SGLANG_QWEN4_PLE_FILE_DIR/<model path>``, one directory per checkpoint."""
+    safe = re.sub(r"[^A-Za-z0-9._-]+", "_", str(model_path).rstrip("/")).strip("_")
+    return os.path.join(envs.SGLANG_QWEN4_PLE_FILE_DIR.get(), safe or "model")
+
+
+def ple_table_file_name(
+    shape: Sequence[int], dtype: torch.dtype, tag: Optional[str] = None
+) -> str:
+    """Deterministic file name so the sparse table is reused across restarts.
+
+    ``tag`` distinguishes tables of the same shape that must not share a file,
+    e.g. the vocabulary shards of different tensor-parallel ranks.
+    """
+    numel = 1
+    for d in shape:
+        numel *= int(d)
+    elem = torch.empty(0, dtype=dtype).element_size()
+    dims = "x".join(str(int(d)) for d in shape)
+    suffix = f"_{tag}" if tag else ""
+    return f"ple_table_{dims}_{str(dtype).replace('torch.', '')}_{numel * elem}B{suffix}.bin"
+
+
+def device_uses_host_page_tables(device_index: int = 0) -> Optional[bool]:
+    """Whether pageable host memory is directly addressable by the GPU.
+
+    Returns None when the CUDA runtime library cannot be queried.
+    """
+    candidates = [ctypes.util.find_library("cudart")]
+    torch_lib = os.path.join(os.path.dirname(torch.__file__), "lib")
+    if os.path.isdir(torch_lib):
+        candidates += sorted(
+            os.path.join(torch_lib, f)
+            for f in os.listdir(torch_lib)
+            if f.startswith("libcudart.so")
+        )
+    try:
+        import nvidia.cuda_runtime  # type: ignore
+
+        nv_lib = os.path.join(os.path.dirname(nvidia.cuda_runtime.__file__), "lib")
+        if os.path.isdir(nv_lib):
+            candidates += sorted(
+                os.path.join(nv_lib, f)
+                for f in os.listdir(nv_lib)
+                if f.startswith("libcudart.so")
+            )
+    except Exception:
+        pass
+    for name in [c for c in candidates if c]:
+        try:
+            cudart = ctypes.CDLL(name)
+            value = ctypes.c_int()
+            rc = cudart.cudaDeviceGetAttribute(
+                ctypes.byref(value),
+                ctypes.c_int(
+                    _CUDA_DEV_ATTR_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES
+                ),
+                ctypes.c_int(device_index),
+            )
+            if rc == 0:
+                return bool(value.value)
+        except OSError:
+            continue
+    return None
+
+
+def _madvise_random(storage: torch.Tensor, nbytes: int) -> None:
+    """The table is pure random access (16 rows of 160 B per token). Without
+    this the kernel's readahead pulls its whole window: measured 1.4 MB of disk
+    per token, ~560x the bytes actually used.
+
+    It bounds readahead I/O only. Folios that are already in the page cache are
+    still mapped in whole on a fault, which is what ``PleFileRssTrimmer``
+    exists for."""
+    if not _madvise(storage.data_ptr(), nbytes, _MADV_RANDOM):
+        logger.warning("PLE table: madvise(MADV_RANDOM) not applied")
+
+
+def _libc() -> Optional[ctypes.CDLL]:
+    global _LIBC
+    if _LIBC is None:
+        try:
+            _LIBC = ctypes.CDLL(
+                ctypes.util.find_library("c") or "libc.so.6", use_errno=True
+            )
+        except OSError:
+            return None
+    return _LIBC
+
+
+def _madvise(addr: int, length: int, advice: int) -> bool:
+    """``madvise(2)`` on our own mapping. Advisory: never affects correctness."""
+    libc = _libc()
+    if libc is None:
+        return False
+    try:
+        rc = libc.madvise(
+            ctypes.c_void_p(addr), ctypes.c_size_t(length), ctypes.c_int(advice)
+        )
+    except Exception:
+        return False
+    if rc != 0:
+        logger.warning(
+            "PLE table: madvise(advice=%d) failed (errno %d)",
+            advice,
+            ctypes.get_errno(),
+        )
+        return False
+    return True
+
+
+def _mapping_rss_bytes(
+    addr: int, nbytes: int, smaps_path: str = "/proc/self/smaps"
+) -> Optional[int]:
+    """Resident bytes of the VMAs overlapping ``[addr, addr + nbytes)``.
+
+    Summed per mapping rather than taken from ``statm``/``smaps_rollup``: only
+    the table's own residency should drive the trim, and on a unified-memory
+    box the process RSS is dominated by everything else.
+    """
+    lo, hi = int(addr), int(addr) + int(nbytes)
+    total = 0
+    overlapping = False
+    try:
+        with open(smaps_path, "r") as f:
+            for line in f:
+                header = _SMAPS_HEADER.match(line)
+                if header is not None:
+                    start = int(header.group(1), 16)
+                    end = int(header.group(2), 16)
+                    overlapping = start < hi and end > lo
+                elif overlapping:
+                    rss = _SMAPS_RSS.match(line)
+                    if rss is not None:
+                        total += int(rss.group(1)) * 1024
+    except OSError:
+        return None
+    return total