[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>
This commit is contained in:
Jimmy Shong
co-authored by
Claude Fable 5.1
rdxa
Yangmin Li
Manrique
yhyang201
Claude Fable 5.1
rdxa
Yangmin Li
Manrique
yhyang201
parent
d6fabb74b4
commit
cebca698e2
@@ -2306,6 +2306,24 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Number of layers per group in offloading.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`-1`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: int</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--ple-offload-embedding`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Keep the Qwen4-Exp PLE n-gram embedding table in host memory instead of device memory; the gather kernel reads it from a host pointer. Enabled by default for BF16 Qwen4-Exp on CUDA; `--no-ple-offload-embedding` disables it.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: bool</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--ple-offload-backend`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Host storage for the offloaded PLE table: `pinned` (CPU pinned memory) or `file` (a sparse file-backed mmap under `--ple-offload-dir`, read directly by the gather kernel). Use `file` on unified-memory devices such as GB10 / DGX Spark, where pinned host memory comes out of the same pool as the weights; it requires a device that reports `cudaDevAttrPageableMemoryAccessUsesHostPageTables`.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`pinned`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: str</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--ple-offload-dir`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Directory for the file-backed PLE table when `--ple-offload-backend` is `file`. The file is sparse and rewritten on every startup; keep it on fast local storage (NVMe). Rewriting an existing table can be much slower than filling a fresh sparse file. Before restarting, stop all servers using that directory and remove its `ple_table_*.bin` files to avoid this startup cost.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`$SGLANG_CACHE_DIR/ple/<model path>`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: str</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--offload-num-in-group`</td>
|
||||
|
||||
@@ -83,6 +83,31 @@ SGLang supports various environment variables that can be used to configure its
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Cache directory for model weights and other data. Also the default root for compiled-kernel caches: Triton, Inductor, FlashInfer, the CUDA driver and DeepGEMM are pointed under it unless their own env vars (`TRITON_CACHE_DIR`, `TORCHINDUCTOR_CACHE_DIR`, `FLASHINFER_WORKSPACE_BASE`, `CUDA_CACHE_PATH`, `SGLANG_DG_CACHE_DIR`) are set explicitly</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>~/.cache/sglang</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_QWEN4_PLE_FILE_DIR</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Default directory for the file-backed Qwen4-Exp PLE table (`--ple-offload-backend file`); one subdirectory per model path is created under it</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>{SGLANG_CACHE_DIR}/ple</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_QWEN4_PLE_FILE_PREFETCH</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Hint the page cache (`posix_fadvise(WILLNEED)`) with the rows a prefill-sized PLE gather is about to read from the file-backed table. Set to `0` to disable</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>1</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_QWEN4_PLE_FILE_SKIP_DEVICE_CHECK</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Skip the load-time check that the device reads pageable host memory through the host page tables (`cudaDevAttrPageableMemoryAccessUsesHostPageTables`) before using the file-backed PLE table; only for devices that cannot be queried</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>0</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_QWEN4_PLE_FILE_RSS_BUDGET_GB</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Cap (GiB) on the resident set of the file-backed PLE mapping. A row fault maps a whole page-cache folio, so the mapping otherwise creeps towards the full table and consumes the free memory that sizes the KV pool; over the cap the pages are dropped with `MADV_DONTNEED` and stay in the page cache. Set to `0` to disable</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>8.0</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_QWEN4_PLE_FILE_RSS_INTERVAL_S</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>How often the file-backed PLE mapping's resident set is compared against `SGLANG_QWEN4_PLE_FILE_RSS_BUDGET_GB`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>30.0</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_CUTE_AOT_CACHE_DIR</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Trusted directory for persistent CuTe DSL AOT objects shared across process restarts. Artifacts are namespaced by source, runtime ABI, host platform, and target GPU architecture. Set to an empty string to keep compilation process-local. Only use a directory writable by trusted users because SGLang loads cached object files into the process.</td>
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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``."""
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -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"
|
||||
):
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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
|
||||
@@ -1,3 +1,5 @@
|
||||
import os
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -79,11 +81,11 @@ def _make_source_embedding(
|
||||
)
|
||||
|
||||
|
||||
def _load_rows(offloaded, rows):
|
||||
def _load_rows(offloaded, rows, *, pinned=True):
|
||||
pointer = offloaded.weight.data_ptr()
|
||||
offloaded.weight_loader(offloaded.weight, rows)
|
||||
assert offloaded.weight.data_ptr() == pointer
|
||||
assert offloaded.weight.is_pinned()
|
||||
assert offloaded.weight.is_pinned() == pinned
|
||||
assert offloaded.weight.weight_loader.__self__ is offloaded
|
||||
assert offloaded.quant_method is None
|
||||
|
||||
@@ -189,6 +191,78 @@ def test_qwen4_ple_prefetch_buffer_lifecycle(monkeypatch):
|
||||
assert set(layer._graph_prefetch_buffers) == {3, 5}
|
||||
|
||||
|
||||
def _file_backend_supported() -> bool:
|
||||
from sglang.srt.models.qwen4_exp_ple_table import device_uses_host_page_tables
|
||||
|
||||
return (
|
||||
torch.cuda.is_available()
|
||||
and device_uses_host_page_tables(torch.cuda.current_device()) is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _file_backend_supported(),
|
||||
reason="the file backend needs pageable host memory reachable through host page tables",
|
||||
)
|
||||
@pytest.mark.parametrize("embedding_dim", [7, 160])
|
||||
def test_qwen4_ple_file_backend_matches_pinned(embedding_dim):
|
||||
with tempfile.TemporaryDirectory() as table_dir:
|
||||
pinned = Qwen4ExpPinnedHostEmbedding(
|
||||
_make_source_embedding(embedding_dim=embedding_dim)
|
||||
)
|
||||
filed = Qwen4ExpPinnedHostEmbedding(
|
||||
_make_source_embedding(embedding_dim=embedding_dim),
|
||||
backend="file",
|
||||
table_dir=table_dir,
|
||||
)
|
||||
assert pinned._file_prefetcher is None and filed._file_prefetcher is not None
|
||||
(name,) = os.listdir(table_dir)
|
||||
assert "rows0-8" in name # this rank's vocabulary shard
|
||||
rows = torch.arange(
|
||||
8 * embedding_dim, dtype=torch.bfloat16, device="cuda"
|
||||
).reshape(8, embedding_dim)
|
||||
_load_rows(pinned, rows)
|
||||
_load_rows(filed, rows, pinned=False)
|
||||
ids = torch.tensor([[0, 7, 3], [4, 1, 6]], dtype=torch.int64, device="cuda")
|
||||
torch.testing.assert_close(filed(ids), pinned(ids), rtol=0, atol=0)
|
||||
# A prefill-sized gather goes through the page-cache hint path.
|
||||
big = torch.randint(0, 8, (4096,), device="cuda")
|
||||
torch.testing.assert_close(
|
||||
filed(big), rows.index_select(0, big), rtol=0, atol=0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not _file_backend_supported(),
|
||||
reason="the file backend needs pageable host memory reachable through host page tables",
|
||||
)
|
||||
def test_qwen4_ple_file_backend_fp8_table():
|
||||
embedding_dim = 160
|
||||
with tempfile.TemporaryDirectory() as table_dir:
|
||||
filed = Qwen4ExpPinnedHostEmbedding(
|
||||
_make_source_embedding(
|
||||
embedding_dim=embedding_dim, dtype=torch.float8_e4m3fn
|
||||
),
|
||||
backend="file",
|
||||
table_dir=table_dir,
|
||||
)
|
||||
assert filed.weight.dtype == torch.float8_e4m3fn
|
||||
rows = (
|
||||
torch.arange(8 * embedding_dim, dtype=torch.float32, device="cuda").reshape(
|
||||
8, embedding_dim
|
||||
)
|
||||
/ 64
|
||||
).to(torch.float8_e4m3fn)
|
||||
_load_rows(filed, rows, pinned=False)
|
||||
ids = torch.tensor([[0, 7, 3]], dtype=torch.int64, device="cuda")
|
||||
expected = (
|
||||
rows.index_select(0, ids.flatten())
|
||||
.to(torch.bfloat16)
|
||||
.reshape(1, 3, embedding_dim)
|
||||
)
|
||||
torch.testing.assert_close(filed(ids), expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
"""Fp8MoEMethod builds a triton runner when the global MoE runner backend is
|
||||
flashinfer_cutlass or flashinfer_cutedsl, which have no fp8 MoE path."""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import sglang.srt.layers.quantization.fp8 as fp8
|
||||
from sglang.srt.layers.moe import MoeRunnerBackend, MoeRunnerConfig
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestFp8MoeRunnerFallback(CustomTestCase):
|
||||
def _runner_backend_for(self, global_backend):
|
||||
method = fp8.Fp8MoEMethod.__new__(fp8.Fp8MoEMethod)
|
||||
with patch.object(fp8, "get_moe_runner_backend", return_value=global_backend):
|
||||
method.create_moe_runner(layer=None, moe_runner_config=MoeRunnerConfig())
|
||||
return method.runner.runner_backend
|
||||
|
||||
def test_flashinfer_cutlass_falls_back_to_triton(self):
|
||||
self.assertTrue(
|
||||
self._runner_backend_for(MoeRunnerBackend.FLASHINFER_CUTLASS).is_triton()
|
||||
)
|
||||
|
||||
def test_flashinfer_cutedsl_falls_back_to_triton(self):
|
||||
self.assertTrue(
|
||||
self._runner_backend_for(MoeRunnerBackend.FLASHINFER_CUTEDSL).is_triton()
|
||||
)
|
||||
|
||||
def test_triton_is_kept(self):
|
||||
self.assertTrue(self._runner_backend_for(MoeRunnerBackend.TRITON).is_triton())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -22,7 +22,12 @@ from sglang.srt.configs.model_config import ModelConfig
|
||||
from sglang.srt.layers.linear import ReplicatedLinear
|
||||
from sglang.srt.layers.logits_processor import should_apply_lm_head_quant_method
|
||||
from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
|
||||
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
|
||||
from sglang.srt.layers.quantization.fp8 import (
|
||||
Fp8Config,
|
||||
Fp8LinearMethod,
|
||||
Fp8MoEMethod,
|
||||
)
|
||||
from sglang.srt.layers.quantization.modelopt_quant import (
|
||||
ModelOptFp4Config,
|
||||
ModelOptFp4LinearMethod,
|
||||
@@ -1188,6 +1193,56 @@ class TestModelOptMixedPrecisionConfig(CustomTestCase):
|
||||
"FP8",
|
||||
)
|
||||
|
||||
def test_mixed_precision_resolves_vl_language_model_keys(self):
|
||||
# nvidia/Qwen3.8-Flash-Next-NVFP4 keys the text stack as
|
||||
# `model.language_model.*` while Qwen4-Exp modules are `model.*`.
|
||||
quant_config = ModelOptMixedPrecisionConfig.from_config(
|
||||
{
|
||||
"quant_algo": "MIXED_PRECISION",
|
||||
"quantized_layers": {
|
||||
"model.language_model.layers.3.mlp.experts": {
|
||||
"quant_algo": "NVFP4",
|
||||
"group_size": 16,
|
||||
},
|
||||
"model.language_model.layers.1.ple.ple_embedding.ngram_embedding": {
|
||||
"quant_algo": "FP8"
|
||||
},
|
||||
"mtp.layers.0.mlp.experts": {
|
||||
"quant_algo": "FP8_BLOCK_SCALES",
|
||||
"group_size": 128,
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
self.assertEqual(quant_config.exclude_modules, [])
|
||||
moe = FusedMoE.__new__(FusedMoE)
|
||||
self.assertIsInstance(
|
||||
quant_config.get_quant_method(moe, "mtp.layers.0.mlp.experts"),
|
||||
Fp8MoEMethod,
|
||||
)
|
||||
self.assertEqual(
|
||||
quant_config.get_quant_method(
|
||||
moe, "mtp.layers.0.mlp.experts"
|
||||
).quant_config.weight_block_size,
|
||||
[128, 128],
|
||||
)
|
||||
self.assertEqual(
|
||||
quant_config.resolve_quant_algo("model.layers.3.mlp.experts"), "NVFP4"
|
||||
)
|
||||
self.assertEqual(
|
||||
quant_config.resolve_quant_algo(
|
||||
"model.layers.1.ple.ple_embedding.ngram_embedding"
|
||||
),
|
||||
"FP8",
|
||||
)
|
||||
self.assertIsNone(
|
||||
quant_config.resolve_quant_algo("model.layers.1.ple.key_proj")
|
||||
)
|
||||
self.assertIsNone(
|
||||
quant_config.resolve_quant_algo("model.layers.3.mlp.shared_expert")
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,332 @@
|
||||
"""File-backed host storage for the offloaded Qwen4-Exp PLE table.
|
||||
|
||||
CPU part: the allocator builds a sparse file of exactly the table's size, hands
|
||||
back a tensor with the requested shape/dtype whose writes land in the file and
|
||||
survive a re-open, reuses the file across calls, replaces one of the wrong size,
|
||||
and the prefetcher computes the right page set and honours its size floor. The
|
||||
resident-set trimmer measures only its own mapping, drops its pages once over
|
||||
budget without losing what was written through them, and is off when the budget
|
||||
is zero or the mapping is pinned.
|
||||
|
||||
GPU part (skipped unless the device reads pageable host memory through the host
|
||||
page tables, i.e. unified-memory parts such as GB10): the production Triton
|
||||
gather kernel reading from the file-backed table matches a torch gather.
|
||||
"""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.models.qwen4_exp_ple_table import (
|
||||
PleFilePrefetcher,
|
||||
PleFileRssTrimmer,
|
||||
_mapping_rss_bytes,
|
||||
allocate_ple_host_table,
|
||||
default_ple_table_dir,
|
||||
device_uses_host_page_tables,
|
||||
make_ple_file_prefetcher,
|
||||
make_ple_file_rss_trimmer,
|
||||
ple_table_file_name,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestPleFileTableAllocator(CustomTestCase):
|
||||
def test_file_is_sparse_and_sized_exactly(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
table = allocate_ple_host_table((1000, 160), torch.float8_e4m3fn, "file", d)
|
||||
path = os.path.join(
|
||||
d, ple_table_file_name((1000, 160), torch.float8_e4m3fn)
|
||||
)
|
||||
self.assertTrue(os.path.exists(path))
|
||||
self.assertEqual(os.path.getsize(path), 1000 * 160)
|
||||
self.assertEqual(tuple(table.shape), (1000, 160))
|
||||
self.assertEqual(table.dtype, torch.float8_e4m3fn)
|
||||
# Sparse: nothing written yet, so (almost) no blocks allocated.
|
||||
self.assertLess(os.stat(path).st_blocks * 512, 64 * 1024)
|
||||
|
||||
def test_writes_persist_and_file_is_reused(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
shape, dtype = (64, 32), torch.bfloat16
|
||||
table = allocate_ple_host_table(shape, dtype, "file", d)
|
||||
row = torch.arange(32, dtype=torch.float32).to(dtype)
|
||||
table[7].copy_(row) # what the weight loader does, row by row
|
||||
del table
|
||||
again = allocate_ple_host_table(shape, dtype, "file", d)
|
||||
self.assertTrue(torch.equal(again[7].float(), row.float()))
|
||||
self.assertEqual(len(os.listdir(d)), 1)
|
||||
|
||||
def test_wrong_sized_file_is_replaced(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
path = os.path.join(d, ple_table_file_name((8, 8), torch.bfloat16))
|
||||
with open(path, "wb") as f:
|
||||
f.write(b"\x01" * 10)
|
||||
table = allocate_ple_host_table((8, 8), torch.bfloat16, "file", d)
|
||||
self.assertEqual(os.path.getsize(path), 8 * 8 * 2)
|
||||
self.assertEqual(tuple(table.shape), (8, 8))
|
||||
|
||||
def test_tag_separates_tensor_parallel_shards(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
a = allocate_ple_host_table(
|
||||
(8, 8), torch.bfloat16, "file", d, tag="rows0-8"
|
||||
)
|
||||
b = allocate_ple_host_table(
|
||||
(8, 8), torch.bfloat16, "file", d, tag="rows8-16"
|
||||
)
|
||||
a.fill_(1.0)
|
||||
self.assertEqual(len(os.listdir(d)), 2)
|
||||
self.assertTrue(torch.all(b.float() == 0.0))
|
||||
self.assertIn(
|
||||
"rows8-16", ple_table_file_name((8, 8), torch.bfloat16, "rows8-16")
|
||||
)
|
||||
|
||||
def test_default_dir_is_per_checkpoint(self):
|
||||
with mock.patch.dict(os.environ, {"SGLANG_QWEN4_PLE_FILE_DIR": "/cache/ple"}):
|
||||
a = default_ple_table_dir("RadixArk/Qwen3.8-Flash-Next-NVFP4")
|
||||
b = default_ple_table_dir("/root/.cache/huggingface/flashnext-fp8/")
|
||||
self.assertEqual(a, "/cache/ple/RadixArk_Qwen3.8-Flash-Next-NVFP4")
|
||||
self.assertEqual(b, "/cache/ple/root_.cache_huggingface_flashnext-fp8")
|
||||
self.assertNotEqual(a, b)
|
||||
|
||||
def test_unknown_backend_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
allocate_ple_host_table((4, 4), torch.bfloat16, "nvme", None)
|
||||
|
||||
def test_pinned_backend_unchanged(self):
|
||||
if not torch.cuda.is_available():
|
||||
self.skipTest("pinned memory needs a CUDA runtime")
|
||||
table = allocate_ple_host_table((4, 4), torch.bfloat16, "pinned", None)
|
||||
self.assertTrue(table.is_pinned())
|
||||
self.assertIsNone(make_ple_file_prefetcher(table))
|
||||
|
||||
|
||||
class TestPleFilePrefetcher(CustomTestCase):
|
||||
def test_page_set_covers_row_start_and_end(self):
|
||||
# 160-byte rows: row 25 spans bytes 4000-4159, i.e. pages 0 and 1.
|
||||
pages = PleFilePrefetcher.pages_for_rows(torch.tensor([25, 0]), 160)
|
||||
self.assertEqual(pages, [0, 1])
|
||||
pages = PleFilePrefetcher.pages_for_rows(torch.tensor([1000, 1000]), 160)
|
||||
self.assertEqual(pages, [39]) # dedup, single page
|
||||
|
||||
def test_enqueue_uses_local_tp_offsets_and_ignores_other_shards(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
path = os.path.join(d, "t.bin")
|
||||
with open(path, "wb") as f:
|
||||
f.truncate(8192)
|
||||
pf = PleFilePrefetcher(path, row_bytes=160, min_rows=1)
|
||||
try:
|
||||
with mock.patch("os.posix_fadvise") as fadvise:
|
||||
self.assertFalse(
|
||||
pf.enqueue(
|
||||
torch.tensor([999, 1032]), vocab_start=1000, vocab_end=1032
|
||||
)
|
||||
)
|
||||
self.assertTrue(
|
||||
pf.enqueue(
|
||||
torch.tensor([999, 1000, 1025, 1032]),
|
||||
vocab_start=1000,
|
||||
vocab_end=1032,
|
||||
)
|
||||
)
|
||||
pf._pool.shutdown(wait=True)
|
||||
self.assertEqual(
|
||||
sorted(c.args[1] for c in fadvise.call_args_list), [0, 4096]
|
||||
)
|
||||
finally:
|
||||
pf.close()
|
||||
|
||||
def test_enqueue_respects_min_rows_and_advises_pages(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
path = os.path.join(d, "t.bin")
|
||||
with open(path, "wb") as f:
|
||||
f.truncate(1 << 20)
|
||||
pf = PleFilePrefetcher(path, row_bytes=160, min_rows=4)
|
||||
try:
|
||||
self.assertFalse(pf.enqueue(torch.tensor([1, 2, 3])))
|
||||
with mock.patch("os.posix_fadvise") as fadvise:
|
||||
self.assertTrue(pf.enqueue(torch.tensor([0, 1, 2, 30])))
|
||||
pf._pool.shutdown(wait=True)
|
||||
offsets = sorted(c.args[1] for c in fadvise.call_args_list)
|
||||
# rows 0-2 live in page 0; row 30 (bytes 4800-4959) in page 1
|
||||
self.assertEqual(offsets, [0, 4096])
|
||||
finally:
|
||||
pf.close()
|
||||
|
||||
|
||||
class TestSmapsParsing(CustomTestCase):
|
||||
"""The parser runs anywhere; only the live mapping needs Linux."""
|
||||
|
||||
SMAPS = """00400000-00401000 r--p 00000000 08:01 1 /usr/bin/x
|
||||
Size: 4 kB
|
||||
Rss: 4 kB
|
||||
7f0000000000-7f0004000000 rw-s 00000000 08:01 2 /cache/ple/table.bin
|
||||
Size: 65536 kB
|
||||
Rss: 32768 kB
|
||||
7f0004000000-7f0008000000 rw-s 00000000 08:01 2 /cache/ple/table.bin
|
||||
Size: 65536 kB
|
||||
Rss: 1024 kB
|
||||
7f0100000000-7f0100001000 rw-p 00000000 00:00 0
|
||||
Size: 4 kB
|
||||
Rss: 4 kB
|
||||
"""
|
||||
|
||||
def _rss(self, addr, nbytes):
|
||||
with tempfile.NamedTemporaryFile("w", suffix=".smaps", delete=False) as f:
|
||||
f.write(self.SMAPS)
|
||||
path = f.name
|
||||
try:
|
||||
return _mapping_rss_bytes(addr, nbytes, smaps_path=path)
|
||||
finally:
|
||||
os.unlink(path)
|
||||
|
||||
def test_sums_every_vma_of_the_table_and_nothing_else(self):
|
||||
# The table spans both of its VMAs; the unrelated ones must not count.
|
||||
self.assertEqual(self._rss(0x7F0000000000, 0x8000000), (32768 + 1024) * 1024)
|
||||
|
||||
def test_counts_a_partially_overlapping_vma(self):
|
||||
# A range ending inside the first VMA still needs that VMA's pages.
|
||||
self.assertEqual(self._rss(0x7F0000000000, 0x1000), 32768 * 1024)
|
||||
|
||||
def test_ignores_unrelated_mappings(self):
|
||||
self.assertEqual(self._rss(0x7F0200000000, 0x1000), 0)
|
||||
|
||||
def test_missing_smaps_is_reported_as_unknown(self):
|
||||
self.assertIsNone(
|
||||
_mapping_rss_bytes(0x1000, 0x1000, smaps_path="/nonexistent/smaps")
|
||||
)
|
||||
|
||||
|
||||
class TestPleFileRssTrimmerConfig(CustomTestCase):
|
||||
def test_budget_zero_disables_the_trimmer(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
table = allocate_ple_host_table((64, 32), torch.bfloat16, "file", d)
|
||||
with mock.patch.dict(
|
||||
os.environ, {"SGLANG_QWEN4_PLE_FILE_RSS_BUDGET_GB": "0"}
|
||||
):
|
||||
self.assertIsNone(make_ple_file_rss_trimmer(table))
|
||||
|
||||
def test_pinned_table_has_no_trimmer(self):
|
||||
if not torch.cuda.is_available():
|
||||
self.skipTest("pinned memory needs a CUDA runtime")
|
||||
table = allocate_ple_host_table((4, 4), torch.bfloat16, "pinned", None)
|
||||
self.assertIsNone(make_ple_file_rss_trimmer(table))
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
os.path.exists("/proc/self/smaps"),
|
||||
"the resident set of a mapping is only readable on Linux",
|
||||
)
|
||||
class TestPleFileRssTrimmer(CustomTestCase):
|
||||
# 64 MiB: large enough that the Rss of the mapping stands out, small
|
||||
# enough to write in a CPU test.
|
||||
SHAPE = (32768, 1024)
|
||||
NBYTES = 32768 * 1024 * 2
|
||||
|
||||
def _trimmer(self, table, budget_bytes):
|
||||
return PleFileRssTrimmer(
|
||||
addr=table.data_ptr(),
|
||||
nbytes=self.NBYTES,
|
||||
budget_bytes=budget_bytes,
|
||||
interval_s=3600.0,
|
||||
chunk_bytes=16 << 20, # several chunks, as in production
|
||||
)
|
||||
|
||||
def test_measures_its_own_mapping_only(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
table = allocate_ple_host_table(self.SHAPE, torch.bfloat16, "file", d)
|
||||
trimmer = self._trimmer(table, 0)
|
||||
empty = trimmer.mapping_rss_bytes()
|
||||
table.fill_(1.0) # touches every page
|
||||
touched = trimmer.mapping_rss_bytes()
|
||||
self.assertIsNotNone(touched)
|
||||
self.assertGreater(touched, empty)
|
||||
self.assertGreater(touched, self.NBYTES // 2)
|
||||
# Never the whole process: only the VMAs backing this table.
|
||||
self.assertLessEqual(touched, self.NBYTES + (16 << 20))
|
||||
|
||||
def test_over_budget_drops_pages_and_keeps_the_data(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
table = allocate_ple_host_table(self.SHAPE, torch.bfloat16, "file", d)
|
||||
table.fill_(1.0)
|
||||
table[7][3] = 2.0
|
||||
trimmer = self._trimmer(table, budget_bytes=1 << 20)
|
||||
before = trimmer.mapping_rss_bytes()
|
||||
freed = trimmer.trim_once()
|
||||
after = trimmer.mapping_rss_bytes()
|
||||
self.assertGreater(freed, 0)
|
||||
self.assertLess(after, before // 2)
|
||||
# MADV_DONTNEED on a shared file mapping drops the page-table
|
||||
# entries, not the page cache: the writes are still there.
|
||||
self.assertEqual(table[7][3].item(), 2.0)
|
||||
self.assertEqual(table[9][3].item(), 1.0)
|
||||
|
||||
def test_under_budget_is_a_no_op(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
table = allocate_ple_host_table(self.SHAPE, torch.bfloat16, "file", d)
|
||||
table.fill_(1.0)
|
||||
trimmer = self._trimmer(table, budget_bytes=self.NBYTES * 4)
|
||||
before = trimmer.mapping_rss_bytes()
|
||||
self.assertEqual(trimmer.trim_once(), 0)
|
||||
self.assertEqual(trimmer.mapping_rss_bytes(), before)
|
||||
|
||||
def test_factory_starts_and_stops_a_thread(self):
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
table = allocate_ple_host_table((64, 32), torch.bfloat16, "file", d)
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"SGLANG_QWEN4_PLE_FILE_RSS_BUDGET_GB": "1",
|
||||
"SGLANG_QWEN4_PLE_FILE_RSS_INTERVAL_S": "3600",
|
||||
},
|
||||
):
|
||||
trimmer = make_ple_file_rss_trimmer(table)
|
||||
self.assertIsNotNone(trimmer)
|
||||
try:
|
||||
self.assertTrue(trimmer._thread.is_alive())
|
||||
finally:
|
||||
trimmer.close()
|
||||
trimmer._thread.join(timeout=5)
|
||||
self.assertFalse(trimmer._thread.is_alive())
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
torch.cuda.is_available() and device_uses_host_page_tables(0) is True,
|
||||
"needs a device that reads pageable host memory through the host page tables",
|
||||
)
|
||||
class TestPleFileTableGatherOnDevice(CustomTestCase):
|
||||
def test_triton_gather_reads_file_backed_table(self):
|
||||
import triton
|
||||
|
||||
from sglang.srt.models.qwen4_exp import (
|
||||
_gather_ple_embedding_from_pinned_kernel,
|
||||
)
|
||||
|
||||
rows, dim = 4096, 160
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
table = allocate_ple_host_table((rows, dim), torch.bfloat16, "file", d)
|
||||
table.copy_(torch.randn(rows, dim).to(torch.bfloat16))
|
||||
ids = torch.randint(0, rows, (2048,), device="cuda")
|
||||
out = torch.empty(2048, dim, dtype=torch.bfloat16, device="cuda")
|
||||
_gather_ple_embedding_from_pinned_kernel[(ids.numel(),)](
|
||||
table.data_ptr(),
|
||||
ids,
|
||||
out,
|
||||
embedding_dim=dim,
|
||||
tp_vocab_start=0,
|
||||
tp_vocab_end=rows,
|
||||
is_fp8=False,
|
||||
BLOCK_D=triton.next_power_of_2(dim),
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
expected = table[ids.cpu()].to("cuda")
|
||||
self.assertTrue(torch.equal(out, expected))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -656,7 +656,20 @@ class TestWrapperEntryClassGates(_FusionGateCase):
|
||||
)
|
||||
|
||||
# The normalization the constructor applies, shared with the gate.
|
||||
self.assertIsNone(_mtp_quant_config(_quant("modelopt_mixed")))
|
||||
mixed_bf16_mtp = SimpleNamespace(
|
||||
get_name=lambda: "modelopt_mixed",
|
||||
quantized_layers={"model.layers.0.mlp.experts": {"quant_algo": "NVFP4"}},
|
||||
)
|
||||
self.assertIsNone(_mtp_quant_config(mixed_bf16_mtp))
|
||||
# MIXED_PRECISION checkpoints that quantize the MTP head keep it.
|
||||
mixed_fp8_mtp = SimpleNamespace(
|
||||
get_name=lambda: "modelopt_mixed",
|
||||
quantized_layers={
|
||||
"model.layers.0.mlp.experts": {"quant_algo": "NVFP4"},
|
||||
"mtp.layers.0.mlp.experts": {"quant_algo": "FP8_BLOCK_SCALES"},
|
||||
},
|
||||
)
|
||||
self.assertIs(_mtp_quant_config(mixed_fp8_mtp), mixed_fp8_mtp)
|
||||
serialized = SimpleNamespace(
|
||||
get_name=lambda: "modelopt_fp4", is_checkpoint_nvfp4_serialized=True
|
||||
)
|
||||
|
||||
@@ -637,6 +637,22 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
pp_size=2,
|
||||
)
|
||||
|
||||
def test_qwen4_ple_file_requires_offload(self):
|
||||
qwen4 = ("Qwen4ExpForConditionalGeneration", "qwen4_exp")
|
||||
with override_platform(is_cuda=True):
|
||||
sa = self._construct(
|
||||
*qwen4,
|
||||
ple_offload_embedding=True,
|
||||
ple_offload_backend="file",
|
||||
ple_offload_dir="/tmp/ple",
|
||||
)
|
||||
self.assertEqual(self._resolved(sa, "ple_offload_backend"), "file")
|
||||
self.assertEqual(self._resolved(sa, "ple_offload_dir"), "/tmp/ple")
|
||||
with self.assertRaisesRegex(ValueError, "requires --ple-offload-embedding"):
|
||||
self._construct(
|
||||
*qwen4, ple_offload_embedding=False, ple_offload_backend="file"
|
||||
)
|
||||
|
||||
def test_qwen4_ple_offload_default(self):
|
||||
qwen4 = ("Qwen4ExpForConditionalGeneration", "qwen4_exp")
|
||||
with override_platform(is_cuda=True):
|
||||
@@ -3072,6 +3088,93 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
with override_platform(is_sm100=False):
|
||||
self.assertEqual(_qwen3_moe_family_overrides(None, None), {})
|
||||
|
||||
def test_qwen3_moe_family_mixed_precision_moe_runner(self):
|
||||
from sglang.srt.arg_groups.model_overrides.qwen3_moe import (
|
||||
_qwen3_moe_family_overrides,
|
||||
)
|
||||
|
||||
def _mixed(expert_algo):
|
||||
return SimpleNamespace(
|
||||
architectures=["Qwen4ExpForConditionalGeneration"],
|
||||
quantization_config={
|
||||
"quant_method": "modelopt_mixed",
|
||||
"quantized_layers": {
|
||||
"model.language_model.layers.0.mlp.experts": {
|
||||
"quant_algo": expert_algo
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
args = SimpleNamespace(
|
||||
quantization="modelopt_mixed",
|
||||
_quantization_explicitly_unset=False,
|
||||
moe_a2a_backend="none",
|
||||
moe_runner_backend="auto",
|
||||
)
|
||||
with override_platform(is_sm100=True):
|
||||
# W4A4 experts take trtllm-gen like modelopt_fp4; W4A16 has no
|
||||
# trtllm-gen kernel and goes to marlin.
|
||||
self.assertEqual(
|
||||
_qwen3_moe_family_overrides(args, _mixed("NVFP4")),
|
||||
{"moe_runner_backend": "flashinfer_trtllm"},
|
||||
)
|
||||
self.assertEqual(
|
||||
_qwen3_moe_family_overrides(args, _mixed("W4A16_NVFP4")),
|
||||
{"moe_runner_backend": "marlin"},
|
||||
)
|
||||
|
||||
def test_qwen3_moe_family_w4a16_explicit_runner(self):
|
||||
"""Keep opted-in CuTe DSL v2 W4A16 accepted and auto routed to Marlin."""
|
||||
from sglang.srt.arg_groups.model_overrides.qwen3_moe import (
|
||||
_qwen3_moe_family_overrides,
|
||||
)
|
||||
|
||||
hf_config = SimpleNamespace(
|
||||
architectures=["Qwen4ExpForConditionalGeneration"],
|
||||
quantization_config={
|
||||
"quant_method": "modelopt_mixed",
|
||||
"quantized_layers": {
|
||||
"model.language_model.layers.0.mlp.experts": {
|
||||
"quant_algo": "W4A16_NVFP4"
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
cases = [
|
||||
("auto", "none", False, {"moe_runner_backend": "marlin"}),
|
||||
("auto", "none", True, {"moe_runner_backend": "marlin"}),
|
||||
("marlin", "none", False, {}),
|
||||
("marlin", "none", True, {}),
|
||||
("flashinfer_cutedsl", "none", True, {}),
|
||||
("flashinfer_cutedsl", "flashinfer", True, {}),
|
||||
("flashinfer_cutedsl", "none", False, None),
|
||||
("flashinfer_cutedsl", "flashinfer", False, None),
|
||||
("flashinfer_cutedsl", "deepep", True, None),
|
||||
("flashinfer_cutlass", "none", True, None),
|
||||
("flashinfer_trtllm", "none", True, None),
|
||||
]
|
||||
for runner, a2a, w4a16_enabled, expected in cases:
|
||||
with (
|
||||
self.subTest(runner=runner, a2a=a2a, w4a16=w4a16_enabled),
|
||||
override_platform(is_sm100=True),
|
||||
envs.SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16.override(w4a16_enabled),
|
||||
):
|
||||
args = SimpleNamespace(
|
||||
quantization=None,
|
||||
_quantization_explicitly_unset=False,
|
||||
moe_a2a_backend=a2a,
|
||||
moe_runner_backend=runner,
|
||||
)
|
||||
if expected is None:
|
||||
with self.assertRaisesRegex(ValueError, "W4A16_NVFP4"):
|
||||
_qwen3_moe_family_overrides(args, hf_config)
|
||||
else:
|
||||
self.assertEqual(
|
||||
_qwen3_moe_family_overrides(args, hf_config),
|
||||
{"quantization": "modelopt_mixed", **expected},
|
||||
)
|
||||
|
||||
def test_step3p_declarations_at_callable_level(self):
|
||||
from sglang.srt.arg_groups.overrides import _step3p_overrides
|
||||
|
||||
|
||||
Reference in New Issue
Block a user