diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 73f899add10f..afa7e7e5d173 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -4093,6 +4093,7 @@ def get_model_loader( model_optloader_allowed = model_config and load_config.load_format not in ( LoadFormat.RUNAI_STREAMER, LoadFormat.REMOTE_INSTANCE, + LoadFormat.IPC_CACHE, ) if model_optloader_allowed and ( diff --git a/python/sglang/srt/weight_cache/daemon.py b/python/sglang/srt/weight_cache/daemon.py index 1c638ddb74ff..8ee5a202c090 100644 --- a/python/sglang/srt/weight_cache/daemon.py +++ b/python/sglang/srt/weight_cache/daemon.py @@ -51,6 +51,7 @@ from .protocol import ( CacheConfig, + capture_module_attrs, check_ipc_quant_support, cleanup_stale_daemon_files, compute_env_stamp, @@ -123,6 +124,13 @@ def __init__( self.config: Optional[CacheConfig] = None # name -> {"handle": base64_str, "shape": list, "dtype": str, "is_param": bool} self.state_entries: Dict[str, Dict[str, Any]] = {} + # module qualname -> {attr: scalar}, the Python-side layout state that + # post-processing stamps and IPC handles cannot carry (see + # capture_module_attrs). + self.module_attrs: Dict[str, Dict[str, Any]] = {} + # MoE runner backend as resolved during this daemon's model load; the + # client compares it once its own model is built. + self.moe_runner_backend: str = "" def _init_distributed(self, server_args, model_config): """Initialize the distributed backend required for model loading. @@ -225,6 +233,14 @@ def load(self): ) publish(server_args, role="weight_cache_daemon") + from sglang.srt.layers.moe import get_moe_runner_backend, initialize_moe_config + from sglang.srt.layers.quantization.fp4_utils import ( + initialize_fp4_gemm_config, + ) + + initialize_fp4_gemm_config(server_args) + initialize_moe_config(server_args) + # Initialize distributed backend for model loading # (must be done after server_args and model_config are available) # Build model config first, then init distributed @@ -269,15 +285,14 @@ def load(self): **compute_env_stamp(), ) - # Refuse to serve quant methods not verified to round-trip through pure - # IPC tensor export. Checked before loading so an unsupported model - # fails fast instead of after minutes of disk I/O. + # Refuse to serve quant methods not verified for IPC sharing. Checked + # before loading so an unsupported model fails fast instead of after + # minutes of disk I/O. check_ipc_quant_support(quant_method, quant_config, where="daemon") # Initialize distributed backend (requires server_args + model_config) self._init_distributed(server_args, model_config) - # Build load config load_config = LoadConfig( load_format=self.load_format, model_loader_extra_config=self.model_loader_extra_config, @@ -290,7 +305,9 @@ def load(self): ) tic = time.perf_counter() - # Load model using DefaultModelLoader (includes TP sharding + quant post-process) + # Load model: the full pipeline, including TP sharding and the quant + # methods' process_weights_after_loading. Clients map the result of that + # pass, so it always runs here and never on a client. loader = get_model_loader(load_config=load_config, model_config=model_config) self.model = loader.load_model( model_config=model_config, @@ -308,6 +325,9 @@ def load(self): # risk observing half-written weights. current_platform.synchronize() + self.module_attrs = capture_module_attrs(self.model, quant_method) + self.moe_runner_backend = get_moe_runner_backend().value + # Export all parameters and buffers as IPC handles self._export_state() @@ -505,6 +525,8 @@ def _handle_connection(self, conn: socket.socket): "status": "ok", "config": self.config.to_dict(), "entries": self.state_entries, + "module_attrs": self.module_attrs, + "moe_runner_backend": self.moe_runner_backend, # PID so the client can watch daemon liveness: if this # process dies while clients hold IPC mappings, their # param.data (and any CUDA-graph-captured addresses) dangle. diff --git a/python/sglang/srt/weight_cache/ipc_loader.py b/python/sglang/srt/weight_cache/ipc_loader.py index 02e7ae772c24..d422794bfcfd 100644 --- a/python/sglang/srt/weight_cache/ipc_loader.py +++ b/python/sglang/srt/weight_cache/ipc_loader.py @@ -26,10 +26,12 @@ from .protocol import ( CacheConfig, + apply_module_attrs, check_ipc_quant_support, compute_env_stamp, get_quant_method_name, hash_quant_config, + ipc_postprocess_reshapes, recv_msg, send_msg, ) @@ -134,13 +136,11 @@ def load_model( device_config, entries, quant_config, + quant_method=quant_method, + module_attrs=cache_data.get("module_attrs") or {}, + daemon_moe_runner_backend=cache_data.get("moe_runner_backend", ""), ) - # Skip _post_load_weights: the daemon already ran - # process_weights_after_loading on the weights before exporting - # IPC handles. Running it again would double-process (e.g., - # re-quantize already-quantized weights), corrupting tensor data. - # Rebuild stale tensor views. Some modules store tensor views as # plain attributes (not parameters/buffers) during __init__. When # the model is initialized on meta device and then weights are @@ -290,6 +290,9 @@ def _load_zero_copy_mode( device_config, entries, quant_config, + quant_method: str = "", + module_attrs: Optional[dict] = None, + daemon_moe_runner_backend: str = "", ) -> nn.Module: """Zero-copy load: map IPC tensors directly as param.data. @@ -297,6 +300,7 @@ def _load_zero_copy_mode( then each parameter's data is replaced with the IPC-mapped GPU tensor. The engine and daemon share the same physical GPU memory via CUDA IPC. """ + from sglang.srt.layers.moe import get_moe_runner_backend from sglang.srt.model_loader.utils import set_default_torch_dtype # Initialize model on meta device to avoid any GPU/CPU memory allocation. @@ -331,8 +335,11 @@ def _load_zero_copy_mode( imported_count = 0 mismatched = [] new_params_count = 0 + reshaped_count = 0 map_tic = time.perf_counter() + postprocess_reshapes = ipc_postprocess_reshapes(quant_method) + # Iterate over ALL daemon entries (not just model params/buffers). # This ensures post-quantization parameters (weight_scale, etc.) # that were created by process_weights_after_loading are also mapped. @@ -350,16 +357,18 @@ def _load_zero_copy_mode( imported_tensor.shape != ref_param.shape or imported_tensor.dtype != ref_param.dtype ): - mismatched.append( - f" {name}: IPC={imported_tensor.shape}/{imported_tensor.dtype} " - f"vs model={ref_param.shape}/{ref_param.dtype}" - ) - del imported_tensor - continue + if not postprocess_reshapes: + mismatched.append( + f" {name}: IPC={imported_tensor.shape}/{imported_tensor.dtype} " + f"vs model={ref_param.shape}/{ref_param.dtype}" + ) + del imported_tensor + continue + reshaped_count += 1 # Replace or register the tensor in the model self._set_module_tensor(model, name, imported_tensor, is_param=is_param) - imported_refs.append(imported_tensor) + imported_refs.append((name, imported_tensor)) imported_count += 1 if name not in existing_names: @@ -376,6 +385,29 @@ def _load_zero_copy_mode( f"not merely uninitialized weights:\n" + "\n".join(mismatched) ) + client_moe_runner_backend = get_moe_runner_backend().value + if ( + daemon_moe_runner_backend + and client_moe_runner_backend + and client_moe_runner_backend != daemon_moe_runner_backend + ): + raise RuntimeError( + f"[IpcModelLoader] MoE runner backend mismatch: the daemon's " + f"weights were post-processed for " + f"'{daemon_moe_runner_backend}' but this engine resolved " + f"'{client_moe_runner_backend}'. The expert weight layout is " + f"backend-specific, so serving these weights would produce " + f"wrong output. Launch both with the same --moe-runner-backend." + ) + + if module_attrs: + changed = apply_module_attrs(model, module_attrs) + logger.info( + f"[IpcModelLoader] Applied daemon module attributes " + f"({changed} changed across {len(module_attrs)} modules)" + ) + self._rebind_quant_state_after_import(model) + # After mapping every daemon entry, any tensor still on the meta device # is one the daemon did NOT provide. Filling it with torch.empty() would # hand the model uninitialized GPU memory — silently producing wrong @@ -412,10 +444,15 @@ def _load_zero_copy_mode( ) map_elapsed = time.perf_counter() - map_tic + if reshaped_count: + logger.info( + f"[IpcModelLoader] Rebound {reshaped_count} tensor(s) to the " + f"daemon's post-processed shape (expected for {quant_method})" + ) # Stash IPC refs on the model to prevent GC (which would unmap the memory) if imported_refs: - model._ipc_imported_tensors = imported_refs + model._ipc_imported_tensors = [tensor for _, tensor in imported_refs] logger.info( f"[IpcModelLoader] Zero-copy: mapped {imported_count} tensors " @@ -424,6 +461,54 @@ def _load_zero_copy_mode( return model + @staticmethod + def _rebind_quant_state_after_import(model) -> None: + """Re-establish quant state that depends on tensor identity. + + Post-processing also wires up Python objects holding references to the + tensors it produced -- NVFP4 MoE hands its input global scale to the + token dispatcher. Those are the daemon's objects, so the client redoes + just that wiring against its own IPC-mapped tensors. + + Kept here rather than as a hook on the quant method so the weight cache + stays self-contained. The cost is that the expression below duplicates + the dispatcher call at the end of + ModelOptNvFp4FusedMoEMethod.process_weights_after_loading; if that call + changes, this must change with it, and only end-to-end output parity + would notice. + """ + from sglang.srt.layers.moe.utils import ( + should_use_flashinfer_cutlass_moe_fp4_allgather, + ) + from sglang.srt.layers.quantization.modelopt_quant import ( + MOE_NVFP4_DISPATCH, + ModelOptNvFp4FusedMoEMethod, + ) + + # Map each quant method to a function that rebinds the quant state. + quant_rebinds = { + ModelOptNvFp4FusedMoEMethod: lambda mod: mod.dispatcher.set_quant_config( + { + "input_global_scale": ( + mod.w13_input_scale_quant + if MOE_NVFP4_DISPATCH + or should_use_flashinfer_cutlass_moe_fp4_allgather() + else None + ) + } + ) + } + + rebound = 0 + for _, module in model.named_modules(): + quant_method = getattr(module, "quant_method", None) + if not quant_method or quant_method not in quant_rebinds: + continue + quant_rebinds[quant_method](module) + rebound += 1 + if rebound: + logger.info(f"[IpcModelLoader] Rebound quant state on {rebound} module(s)") + def _fetch_from_cache(self, model_config) -> Optional[dict]: """Connect to daemon, validate config, fetch IPC handles. diff --git a/python/sglang/srt/weight_cache/protocol.py b/python/sglang/srt/weight_cache/protocol.py index 8b3b4b3411cf..69323c9a9724 100644 --- a/python/sglang/srt/weight_cache/protocol.py +++ b/python/sglang/srt/weight_cache/protocol.py @@ -52,6 +52,7 @@ class CacheConfig(msgspec.Struct): # Comparing these turns that into a clean mismatch. See compute_env_stamp(). device_capability: str # local compute capability, e.g. "8.0" ("" if N/A) torch_version: str # torch.__version__ of the process that built the weights + fp4_gemm_backend: str def matches(self, other: "CacheConfig") -> bool: """Check if two configs are compatible for weight sharing.""" @@ -155,13 +156,37 @@ def _fp8_round_trips_via_ipc(quant_config: Any) -> bool: return _get_quant_field(quant_config, "weight_block_size") is not None +def _nvfp4_round_trips_via_ipc(quant_config: Any) -> bool: + """Only plain serialized NVFP4 is verified.""" + quant_algo = _get_quant_field(quant_config, "quant_algo") + if quant_algo is None: + quantization_section = _get_quant_field(quant_config, "quantization") + if isinstance(quantization_section, dict): + quant_algo = quantization_section.get("quant_algo", None) + return quant_algo == "NVFP4" + + # quant_method name -> predicate(quant_config) -> bool (True == verified safe). # A method absent from this registry is unsupported and hard-errors. IPC_QUANT_ALLOWLIST = { "": lambda _quant_config: True, # unquantized "fp8": _fp8_round_trips_via_ipc, # only block-wise FP8 verified + "modelopt_fp4": _nvfp4_round_trips_via_ipc, # NVFP4 } +# quant methods whose post-processing changes tensor shapes (padding for kernel +# alignment, swizzled scale layouts). Their exported tensors legitimately do not +# match the shapes the client's create_weights registered, so the client rebinds +# to the daemon's shape instead of hard-erroring on the mismatch. Correctness +# still rests on the CacheConfig fingerprint (quant config, resolved backends, +# device capability, torch version) matching. +IPC_POSTPROCESS_RESHAPES_QUANTS = {"modelopt_fp4"} + + +def ipc_postprocess_reshapes(quant_method: str) -> bool: + """Whether this method's post-processing changes exported tensor shapes.""" + return quant_method in IPC_POSTPROCESS_RESHAPES_QUANTS + def is_ipc_quant_supported(quant_method: str, quant_config: Any) -> bool: """Return True if `quant_method` is verified safe for IPC zero-copy sharing.""" @@ -240,15 +265,17 @@ def _recv_exact(sock, n: int) -> Optional[bytes]: def compute_env_stamp() -> Dict[str, str]: """Local environment fingerprint for the IPC weight cache. - Returns the device compute capability and torch version of the current - process. A daemon and a connecting client that differ on either may have run - different post-processing / kernel-selection branches, producing weights that - map cleanly over IPC yet serve garbage; stamping these into CacheConfig turns - that into a clean mismatch. Imported lazily so protocol.py stays cheap to - import and usable on CPU-only hosts (both fields degrade to ""). + Returns the device compute capability, torch version, and resolved FP4 GEMM + backend of the current process. A daemon and a connecting client that differ + on any of these may have run different post-processing / kernel-selection + branches (or, for FP4, registered a different raw param set), producing + weights that map cleanly over IPC yet serve garbage; stamping these into + CacheConfig turns that into a clean mismatch. Imported lazily so protocol.py + stays cheap to import and usable on CPU-only hosts (all fields degrade to ""). """ device_capability = "" torch_version = "" + fp4_gemm_backend = "" try: import torch @@ -263,7 +290,101 @@ def compute_env_stamp() -> Dict[str, str]: device_capability = f"{cap.major}.{cap.minor}" except Exception: pass - return {"device_capability": device_capability, "torch_version": torch_version} + try: + from sglang.srt.layers.quantization.fp4_utils import ( + get_fp4_gemm_runner_backend, + ) + + fp4_gemm_backend = str(get_fp4_gemm_runner_backend().value) + except Exception: + pass + return { + "device_capability": device_capability, + "torch_version": torch_version, + "fp4_gemm_backend": fp4_gemm_backend, + } + + +_TRANSFERABLE_ATTR_NAMES = { + # Unquantized: post-processing stamps no layout state. + "": frozenset(), + "fp8": frozenset(), + "modelopt_fp4": frozenset( + { + "weights_padding_cols", + "output_size_per_partition", + "intermediate_size_per_partition", + "_w13_deinterleaved", + } + ), +} +_TRANSFERABLE_ATTR_TYPES = (int, float, bool, str, type(None)) + + +def capture_module_attrs(model, quant_method: str) -> Dict[str, Dict[str, Any]]: + """Snapshot the layout attributes this quant method stamps (see + _TRANSFERABLE_ATTR_NAMES). + + process_weights_after_loading does not only rewrite tensors: it stamps plain + Python attributes that the kernels read later. Tensors ride the IPC handles; + these do not, and a client left with its pre-post-process values pads + activations wrong and serves garbage. So the daemon captures them after + post-processing and the client stamps them on. + + Captured by name rather than as a before/after diff: a diff would silently + miss an attribute that post-processing sets to the value it already had, and + the client applies these onto an identically-configured model, so an + unchanged value is a harmless no-op. + """ + names = _TRANSFERABLE_ATTR_NAMES.get(quant_method, frozenset()) + if not names: + return {} + + attrs: Dict[str, Dict[str, Any]] = {} + for name, module in model.named_modules(): + captured = { + key: value + for key, value in module.__dict__.items() + if key in names and isinstance(value, _TRANSFERABLE_ATTR_TYPES) + } + if captured: + attrs[name] = captured + return attrs + + +def apply_module_attrs(model, module_attrs: Dict[str, Dict[str, Any]]) -> int: + """Stamp daemon-captured plain attributes onto the client's modules. + + Returns the number of attributes whose value actually changed. Raises if the + daemon captured a module the client does not have: that means the two + processes built structurally different models, and silently skipping the + stamp would leave the client reading pre-post-process layout values. + """ + modules = dict(model.named_modules()) + missing = [name for name in module_attrs if name not in modules] + if missing: + raise RuntimeError( + f"[weight_cache] The daemon captured attributes for " + f"{len(missing)} module(s) that do not exist on the client, so the " + f"two processes built structurally different models: " + f"{missing[:10]}{'...' if len(missing) > 10 else ''}" + ) + + changed = 0 + for module_name, captured in module_attrs.items(): + module = modules[module_name] + for key, value in captured.items(): + previous = getattr(module, key, _ATTR_UNSET) + # Only compare like with like: an absent attribute (sentinel) or a + # client-side tensor under the same name must count as changed + # rather than go through a tensor `!=`, which returns a tensor. + if not isinstance(previous, _TRANSFERABLE_ATTR_TYPES) or previous != value: + changed += 1 + setattr(module, key, value) + return changed + + +_ATTR_UNSET = object() def compute_global_rank(tp_size: int, pp_rank: int, tp_rank: int) -> int: diff --git a/test/registered/model_loading/test_weight_cache_daemon.py b/test/registered/model_loading/test_weight_cache_daemon.py index a49dfc4461cc..a4a2f69689a8 100644 --- a/test/registered/model_loading/test_weight_cache_daemon.py +++ b/test/registered/model_loading/test_weight_cache_daemon.py @@ -8,6 +8,11 @@ import torch from sglang.srt.utils import kill_process_tree +from sglang.srt.weight_cache.protocol import ( + compute_global_rank, + get_ready_path, + get_socket_path, +) from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, @@ -43,6 +48,78 @@ ] +def _rank_paths(tp_size: int, pp_rank: int = 0): + """Ready/socket paths for every TP rank, via the protocol's own formula.""" + for tp_rank in range(tp_size): + global_rank = compute_global_rank( + tp_size=tp_size, pp_rank=pp_rank, tp_rank=tp_rank + ) + yield global_rank, get_ready_path(global_rank), get_socket_path(global_rank) + + +def cleanup_daemon_files(tp_size: int) -> None: + """Remove ready/socket files left behind by an earlier (crashed) run. + + A stale .ready would make wait_for_daemon_ready return before the daemon has + loaded anything, so the client would then connect to a socket nobody serves. + """ + for _, ready_path, socket_path in _rank_paths(tp_size): + for path in (ready_path, socket_path): + if os.path.exists(path): + try: + os.unlink(path) + except OSError: + pass + + +def launch_weight_cache_daemon( + model: str, tp_size: int, extra_args=() +) -> subprocess.Popen: + """Start the weight cache daemon covering every TP rank.""" + return subprocess.Popen( + [ + sys.executable, + "-m", + "sglang.srt.weight_cache.daemon", + "--model-path", + model, + "--tp-size", + str(tp_size), + *extra_args, + ] + ) + + +def wait_for_daemon_ready( + daemon_process: subprocess.Popen, + tp_size: int, + timeout: float = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, +) -> None: + """Block until every rank has written its ready file. + + Polls the daemon's exit status too: a daemon that died during loading would + otherwise keep the caller waiting out the full timeout for a file that will + never appear. On timeout the daemon is killed here — callers hit this from + setUpClass, and unittest skips tearDownClass when setUpClass raises, so the + process would leak otherwise. + """ + start = time.time() + for global_rank, ready_path, _ in _rank_paths(tp_size): + while not os.path.exists(ready_path): + if daemon_process.poll() is not None: + raise RuntimeError( + f"Weight cache daemon exited prematurely with code " + f"{daemon_process.returncode}" + ) + if time.time() - start > timeout: + kill_process_tree(daemon_process.pid) + raise TimeoutError( + f"Weight cache daemon rank {global_rank} not ready " + f"within {timeout}s" + ) + time.sleep(2) + + @unittest.skipIf( torch.cuda.device_count() < 2, "TP=2 weight cache daemon test requires >=2 GPUs (skipped on the 1-gpu runner)", @@ -56,44 +133,14 @@ def setUpClass(cls): cls.base_url = DEFAULT_URL_FOR_TEST cls.tp_size = 2 - # Clean up stale ready/socket files from previous runs - for rank in range(cls.tp_size): - for suffix in (".ready", ".sock"): - path = f"/tmp/sglang_weight_cache_rank{rank}{suffix}" - if os.path.exists(path): - os.unlink(path) + cleanup_daemon_files(cls.tp_size) # Step 1: Launch weight cache daemons (blocks until all ranks are ready, # then monitors child processes) - cls.daemon_process = subprocess.Popen( - [ - sys.executable, - "-m", - "sglang.srt.weight_cache.daemon", - "--model-path", - cls.model, - "--tp-size", - str(cls.tp_size), - ] - ) + cls.daemon_process = launch_weight_cache_daemon(cls.model, cls.tp_size) # Step 2: Wait for all daemon ready files - timeout = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH - start = time.time() - for rank in range(cls.tp_size): - ready_path = f"/tmp/sglang_weight_cache_rank{rank}.ready" - while not os.path.exists(ready_path): - if time.time() - start > timeout: - kill_process_tree(cls.daemon_process.pid) - raise TimeoutError( - f"Weight cache daemon rank {rank} not ready within {timeout}s" - ) - if cls.daemon_process.poll() is not None: - raise RuntimeError( - f"Weight cache daemon exited prematurely " - f"with code {cls.daemon_process.returncode}" - ) - time.sleep(2) + wait_for_daemon_ready(cls.daemon_process, cls.tp_size) # Step 3: Launch server in client mode — loads weights via IPC from daemons cls.stdout = open(STDOUT_FILENAME, "w") @@ -129,14 +176,7 @@ def tearDownClass(cls): os.unlink(path) except OSError: pass - for rank in range(getattr(cls, "tp_size", 2)): - for suffix in (".ready", ".sock"): - path = f"/tmp/sglang_weight_cache_rank{rank}{suffix}" - if os.path.exists(path): - try: - os.unlink(path) - except OSError: - pass + cleanup_daemon_files(getattr(cls, "tp_size", 2)) def test_generate(self): for prompt in PROMPTS: @@ -214,44 +254,14 @@ def setUpClass(cls): cls.base_url = DEFAULT_URL_FOR_TEST cls.tp_size = 1 - # Clean up stale ready/socket files from previous runs. - for rank in range(cls.tp_size): - for suffix in (".ready", ".sock"): - path = f"/tmp/sglang_weight_cache_rank{rank}{suffix}" - if os.path.exists(path): - os.unlink(path) + cleanup_daemon_files(cls.tp_size) # Step 1: Launch the weight cache daemon (blocks until the rank is # ready, then monitors the child process). - cls.daemon_process = subprocess.Popen( - [ - sys.executable, - "-m", - "sglang.srt.weight_cache.daemon", - "--model-path", - cls.model, - "--tp-size", - str(cls.tp_size), - ] - ) + cls.daemon_process = launch_weight_cache_daemon(cls.model, cls.tp_size) # Step 2: Wait for the daemon ready file. - timeout = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH - start = time.time() - for rank in range(cls.tp_size): - ready_path = f"/tmp/sglang_weight_cache_rank{rank}.ready" - while not os.path.exists(ready_path): - if time.time() - start > timeout: - kill_process_tree(cls.daemon_process.pid) - raise TimeoutError( - f"Weight cache daemon rank {rank} not ready within {timeout}s" - ) - if cls.daemon_process.poll() is not None: - raise RuntimeError( - f"Weight cache daemon exited prematurely " - f"with code {cls.daemon_process.returncode}" - ) - time.sleep(2) + wait_for_daemon_ready(cls.daemon_process, cls.tp_size) # Step 3: Launch server in client mode — loads weights via IPC. cls.stdout = open(SMOKE_STDOUT_FILENAME, "w") @@ -287,14 +297,7 @@ def tearDownClass(cls): os.unlink(path) except OSError: pass - for rank in range(getattr(cls, "tp_size", 1)): - for suffix in (".ready", ".sock"): - path = f"/tmp/sglang_weight_cache_rank{rank}{suffix}" - if os.path.exists(path): - try: - os.unlink(path) - except OSError: - pass + cleanup_daemon_files(getattr(cls, "tp_size", 1)) def test_generate(self): resp = requests.post( diff --git a/test/registered/model_loading/test_weight_cache_nvfp4.py b/test/registered/model_loading/test_weight_cache_nvfp4.py new file mode 100644 index 000000000000..324d3b40e0da --- /dev/null +++ b/test/registered/model_loading/test_weight_cache_nvfp4.py @@ -0,0 +1,198 @@ +import os +import unittest + +import requests +from test_weight_cache_daemon import ( + cleanup_daemon_files, + launch_weight_cache_daemon, + wait_for_daemon_ready, +) + +from sglang.srt.utils import get_device_sm, kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, + try_cached_model, +) + +MOE_MODEL = DEFAULT_MODEL_NAME_FOR_TEST_MOE_NVFP4 + +# The first prompt doubles as the output sanity check: greedy decoding of a fact +# this small must contain the answer, so grossly wrong weights show up as a +# failure rather than as a 200 with fluent nonsense. +SANITY_PROMPT = "The capital of France is" +SANITY_EXPECTED = "Paris" + +PROMPTS = [ + SANITY_PROMPT, + "The first three prime numbers are", + "Water is made of hydrogen and", +] + +# Two classes (TP=1, TP=2), each costing a daemon disk load plus a client IPC +# map (the map itself is ~0.3s). The TP=2 class self-skips below 2 GPUs. +register_cuda_ci(est_time=700, stage="base-c", runner_config="4-gpu-b200") + + +def _greedy_outputs(base_url, model): + """Deterministic (temperature 0) completions for the fixed prompt set.""" + outputs = [] + for prompt in PROMPTS: + resp = requests.post( + f"{base_url}/v1/completions", + json={ + "model": model, + "prompt": prompt, + "max_tokens": 32, + "temperature": 0, + }, + ) + assert resp.status_code == 200, resp.text + outputs.append(resp.json()["choices"][0]["text"]) + return outputs + + +class NvFp4WeightCacheBase: + """Launch a daemon, then client(s) that map its post-processed weights. + + A plain mixin (not a TestCase) so it is not collected on its own; concrete + classes inherit ``(NvFp4WeightCacheBase, CustomTestCase)`` and set + ``model_path`` / ``check_no_corruption``. + """ + + model_path: str = None + tp_size: int = 1 + # A second client (from the same daemon) costs another server launch, so it + # is opt-in per class. + check_no_corruption: bool = False + + @classmethod + def setUpClass(cls): + cls.model = try_cached_model(cls.model_path) + cls.base_url = DEFAULT_URL_FOR_TEST + cls.daemon_process = None + # (path, handle) per launched client; closed in tearDownClass, see + # _run_client for why they cannot be closed any earlier. + cls._client_logs = [] + + common_args = ["--trust-remote-code", "--quantization", "modelopt_fp4"] + + cleanup_daemon_files(cls.tp_size) + + cls.daemon_process = launch_weight_cache_daemon( + cls.model, cls.tp_size, extra_args=common_args + ) + wait_for_daemon_ready(cls.daemon_process, cls.tp_size) + + cls.client_outputs, cls.client_log_text = cls._run_client(common_args) + + cls.client2_outputs = None + if cls.check_no_corruption: + cls.client2_outputs, _ = cls._run_client(common_args) + + @classmethod + def _run_client(cls, common_args): + """Launch a client-mode server, collect greedy outputs + its logs, kill it. + + The log handle deliberately outlives this call. popen_launch_server tails + the server's stdout/stderr from background threads that run until the + child exits, so closing the sink here -- while the server is still up -- + makes every later line raise "I/O operation on closed file" in those + threads. tearDownClass closes them instead, well after the kill. + """ + log_path = ( + f"/tmp/test_weight_cache_nvfp4_client{len(cls._client_logs)}" + f"_{os.getpid()}.log" + ) + log = open(log_path, "w") + cls._client_logs.append((log_path, log)) + + proc = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + *common_args, + "--tp", + str(cls.tp_size), + "--weight-cache-mode", + "client", + ], + return_stdout_stderr=(log, log), + ) + try: + outputs = _greedy_outputs(cls.base_url, cls.model) + finally: + kill_process_tree(proc.pid) + + log.flush() + log_text = "" + if os.path.exists(log_path): + with open(log_path, errors="replace") as f: + log_text = f.read() + return outputs, log_text + + @classmethod + def tearDownClass(cls): + if getattr(cls, "daemon_process", None) is not None: + kill_process_tree(cls.daemon_process.pid) + cleanup_daemon_files(cls.tp_size) + for log_path, log in getattr(cls, "_client_logs", []): + try: + log.close() + except OSError: + pass + if os.path.exists(log_path): + try: + os.unlink(log_path) + except OSError: + pass + + def test_loaded_via_ipc(self): + self.assertIn( + "Loaded model via IPC", + self.client_log_text, + "client did not load via IPC (likely fell back to disk)", + ) + self.assertIn( + "Applied daemon module attributes", + self.client_log_text, + "client did not re-apply the daemon's post-processing layout state", + ) + + def test_output_is_sane(self): + for prompt, text in zip(PROMPTS, self.client_outputs): + self.assertGreater(len(text.strip()), 0, f"empty output for: {prompt}") + self.assertIn( + SANITY_EXPECTED, + self.client_outputs[0], + f"greedy completion of {SANITY_PROMPT!r} did not contain " + f"{SANITY_EXPECTED!r} -- the IPC-mapped weights are likely wrong", + ) + + def test_no_corruption_second_client(self): + if not self.check_no_corruption: + self.skipTest("no-corruption check is opt-in per class") + self.assertEqual( + self.client2_outputs, + self.client_outputs, + "second client diverged from the first -- the daemon's exported " + "weights were mutated while serving client #1", + ) + + +@unittest.skipIf(get_device_sm() < 100, "NVFP4 requires CUDA SM 100 or higher") +class TestNvFp4WeightCacheMoE(NvFp4WeightCacheBase, CustomTestCase): + model_path = MOE_MODEL + + # Enable tp_size > 1 or second client leads to longer test duration + tp_size = 1 + check_no_corruption = False + + +if __name__ == "__main__": + unittest.main(verbosity=3) diff --git a/test/registered/unit/model_loader/test_weight_cache_protocol.py b/test/registered/unit/model_loader/test_weight_cache_protocol.py index 4788897fdb89..d1779f03632a 100644 --- a/test/registered/unit/model_loader/test_weight_cache_protocol.py +++ b/test/registered/unit/model_loader/test_weight_cache_protocol.py @@ -23,7 +23,6 @@ import unittest from sglang.srt.weight_cache.protocol import ( - IPC_QUANT_ALLOWLIST, CacheConfig, UnsupportedQuantForIPCError, check_ipc_quant_support, @@ -34,6 +33,7 @@ get_ready_path, get_socket_path, hash_quant_config, + ipc_postprocess_reshapes, is_ipc_quant_supported, recv_msg, send_msg, @@ -60,6 +60,7 @@ def _make_cache_config(**overrides) -> CacheConfig: revision="", device_capability="8.0", torch_version="2.5.1", + fp4_gemm_backend="", ) base.update(overrides) return CacheConfig(**base) @@ -126,6 +127,9 @@ def test_any_field_difference_breaks_match(self): ("revision", "v2"), ("device_capability", "9.0"), ("torch_version", "2.4.0"), + # Differing FP4 backends change the exported param set, so this + # must be a clean mismatch rather than a silent wrong-layout load. + ("fp4_gemm_backend", "flashinfer_trtllm"), ): self.assertFalse( base.matches(_make_cache_config(**{field: value})), @@ -224,9 +228,41 @@ def test_block_fp8_supported_but_per_tensor_fp8_rejected(self): self.assertFalse(is_ipc_quant_supported("fp8", {})) self.assertFalse(is_ipc_quant_supported("fp8", None)) + def test_nvfp4_supported(self): + # Accepted in both config layouts: quant_algo at the top level + # (config.json's quantization_config) and nested under "quantization" + # (hf_quant_config.json). ModelOptFp4Config reads both, so both must + # resolve here or a valid checkpoint is refused the zero-copy path. + self.assertTrue(is_ipc_quant_supported("modelopt_fp4", {"quant_algo": "NVFP4"})) + self.assertTrue( + is_ipc_quant_supported( + "modelopt_fp4", {"quantization": {"quant_algo": "NVFP4"}} + ) + ) + + def test_non_nvfp4_modelopt_variants_rejected(self): + # "modelopt_fp4" is a method name, not a format: ModelOptFp4Config also + # accepts FP8 (not the packed NVFP4 layout the daemon exports) and + # NVFP4_AWQ (pre-quant scales, own post-processing). Neither was + # verified over IPC, so the method name alone must not admit them. + self.assertFalse(is_ipc_quant_supported("modelopt_fp4", {"quant_algo": "FP8"})) + self.assertFalse( + is_ipc_quant_supported("modelopt_fp4", {"quant_algo": "NVFP4_AWQ"}) + ) + self.assertFalse( + is_ipc_quant_supported( + "modelopt_fp4", {"quantization": {"quant_algo": "NVFP4_AWQ"}} + ) + ) + # No quantization_config at all: nothing states the checkpoint is NVFP4. + self.assertFalse(is_ipc_quant_supported("modelopt_fp4", None)) + def test_unknown_method_rejected(self): self.assertFalse(is_ipc_quant_supported("gptq_marlin", None)) self.assertFalse(is_ipc_quant_supported("awq", None)) + # Quantize-on-load NVFP4 reports a different method name and must NOT be + # picked up by the modelopt_fp4 (serialized) allowlist entry. + self.assertFalse(is_ipc_quant_supported("nvfp4_online", None)) def test_check_raises_on_unsupported(self): with self.assertRaises(UnsupportedQuantForIPCError): @@ -242,9 +278,10 @@ def test_check_passes_on_supported(self): "fp8", {"weight_block_size": [128, 128]}, where="daemon" ) - def test_allowlist_registry_shape(self): - # Guard against accidentally widening the allowlist without review. - self.assertEqual(set(IPC_QUANT_ALLOWLIST), {"", "fp8"}) + def test_only_reshaping_methods_relax_the_shape_check(self): + self.assertTrue(ipc_postprocess_reshapes("modelopt_fp4")) + self.assertFalse(ipc_postprocess_reshapes("fp8")) + self.assertFalse(ipc_postprocess_reshapes("")) class TestCleanupStaleDaemonFiles(CustomTestCase):