Cheng Wan
b99175dc7d
Last of four; stacked on #38048. The record is the operator's input; the bags are what is in effect. A reader that takes the record and reads a field off it gets the input, which is the wrong one of the two whenever resolution decided something -- and the mistake is silent, because for most fields and most launches the two agree. Several of these files already read both ways, sometimes in the same expression: ```python get_tokenizer( get_serving().tokenizer_path, tokenizer_mode=server_args.tokenizer_mode, # the input, not the decision ... ) ``` Sixty-odd files convert. Record field reads in runtime code go from 199 to 11. Nine parameters that the conversion emptied are dropped along with the argument at every call site -- the dead-parameter ratchet is what names them. ### "Runs after its process publishes" is a per-entry-point claim Most converted reads sit in the serving and model-executor layers, which only exist after publication, or in the two subprocess entry points, which publish first thing. Three places are not like that, and they keep reading the record they were handed: - **`HttpServerEngineAdapter`** launches the server as a *child*. The parent resolves the record and never publishes, so the adapter's own reads -- the launch banner, the API key in its readiness loop, the TP width in `update_weights_from_tensor` -- are of `self.server_args`. A bag read here fails closed in a bare process, or answers for an unrelated engine in one that happens to have published. - **`serve_grpc`** reads its sidecar port before the integrated servicer builds the `Engine` that publishes. The comment above that line already said so and already bound `cfg = resolving_view(server_args)` for it; the sidecar port and the port it derives from read `cfg`. - **`initialize_dp_attention`** runs from callers whose publish is not guaranteed, so its one predicate stays on the resolution view. `ROLE_NAMESPACE_SETS["dp_controller"]` gains `observability` and `serving`, because the controller's metrics gate, tracing setup and worker-port broadcast now read those namespaces. Under `SGLANG_ROLE_NAMESPACES=enforce` that set is what the process may read, so a conversion that reaches a new namespace has to widen it in the same change. ## Three things worth a reviewer's attention **Eleven reads were `getattr(record, "field", default)`.** An AST scan for attribute access does not see those, so the census that said "43 readers" was counting the shape it could match rather than the thing it was after. `incremental_streaming_output` was read that way twice, and the transcription tests were the only reason it surfaced. **Not every record read is a bag read waiting to happen.** A multimodal processor's `base_gpu_id` is the instance's, not the process's: two engines in one process keep different ones, and `test_publishing_another_config_does_not_move_the_device` exists to say so. It stays on the record while `rl_on_policy_target` beside it moves. `RequestMetricsExporter` is the same shape -- it is handed the directory it writes to, and a test builds several with different ones. `configure_logger` is a third: 17 call sites, one of which passes an `argparse.Namespace`, so it is not a global-context reader at all. Those eleven remaining reads are the ones with a reason. **The fixtures move with the code.** Tests that hung config off a mock manager now publish a record, which is what the serving layer reads; where a test states a value it says so with `override_server_args` instead of assigning through the mock. `test_hisparse_unit` is the last of them: it stubbed a `server_args` onto a fake scheduler to say the decode radix cache was off, and the value it was standing in for is the published default, so the stub goes and the class publishes. ## Two things CI caught that a local sweep could not **`unittest.TestCase.enterContext` is Python 3.11+.** The converted fixtures used it at 18 sites; `requires-python` is `>=3.10` and CI runs 3.10, so every one of them raised `AttributeError` there while passing on a newer local interpreter. They call `enter_override(self, ...)` now -- a four-line helper in `sglang/test/test_utils.py` over the override's own `install()` / `restore()`. **A batched sweep cannot see a missing publish.** Three fixtures needed a published config and did not have one; each *passed* inside a shard where some other file had published, and failed when run alone. The affected cases are `test_serving_completions` (which set `incremental_streaming_output` on the mock manager's record, where nothing reads it now), `test_qwen3_vl_feature_materialization` (same shape for `mm_enable_dp_encoder`), and the two Qwen Rust tests -- whose fixture already carried the comment `# Non-auto: get_resolved_model_impl would choke on a SimpleNamespace` next to the `model_impl` it sets, which is exactly what happened once `get_mm_processor_cls` started reading that value from the bag. Its `publish` mirrors `model_impl` now, like the four fields it already mirrored. ## Verification A full registered-unit sweep (648 files) against this stack's merge-base: 19 failures on both sides, the same 19, none of them config. That sweep is what caught 23 failures the file-scoped runs missed -- and, later, that the narrower 139-file list did not even contain the files this change reaches. It is also what caught the `test_hisparse_unit` fixture above: the file passes inside a shard where something else published, and fails when it is run on its own, which is why every failing file is re-run alone before it is counted.
134 lines
5.2 KiB
Python
134 lines
5.2 KiB
Python
"""Unit tests for srt/layers/layernorm_sp (Megatron LayerNorm sequence parallelism).
|
|
|
|
Covers the pure logic that gates SP -- the Qwen3 allowlist, the config guards, and
|
|
the prefill-only activation rule -- without launching a server. The collectives and
|
|
fused matmul fast-paths need a real TP group and are covered by the e2e test.
|
|
"""
|
|
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
|
|
from sglang.srt.arg_groups.layernorm_sp_hook import validate_layernorm_sp
|
|
from sglang.srt.layers import layernorm_sp
|
|
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
|
from sglang.srt.runtime_context import (
|
|
get_flags,
|
|
get_forward,
|
|
publish,
|
|
reset_context,
|
|
)
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cpu_ci(est_time=9, suite="base-a-test-cpu")
|
|
|
|
|
|
def _initialize(*, enable=True, arch="Qwen3ForCausalLM"):
|
|
publish(
|
|
ServerArgs(model_path="dummy", enable_layernorm_sp=enable), role="tokenizer"
|
|
)
|
|
layernorm_sp.initialize_layernorm_sp(
|
|
model_config=SimpleNamespace(
|
|
hf_config=SimpleNamespace(architectures=[arch] if arch else [])
|
|
),
|
|
)
|
|
|
|
|
|
class TestLayerNormSPGating(CustomTestCase):
|
|
def tearDown(self):
|
|
reset_context()
|
|
|
|
def test_initialize_enables_only_for_allowlisted_arch(self):
|
|
_initialize(enable=True, arch="Qwen3ForCausalLM")
|
|
self.assertTrue(layernorm_sp.layernorm_sp_enabled())
|
|
# Flag on but unsupported architecture -> stays off.
|
|
_initialize(enable=True, arch="LlamaForCausalLM")
|
|
self.assertFalse(layernorm_sp.layernorm_sp_enabled())
|
|
# Supported architecture but flag off -> stays off.
|
|
_initialize(enable=False, arch="Qwen3ForCausalLM")
|
|
self.assertFalse(layernorm_sp.layernorm_sp_enabled())
|
|
|
|
def test_defaults_off_before_initialization(self):
|
|
# A process that never runs initialize_layernorm_sp must not enable SP.
|
|
self.assertFalse(layernorm_sp.layernorm_sp_enabled())
|
|
|
|
def test_runs_sp_only_on_extend(self):
|
|
with get_flags().sp.override(enabled=True):
|
|
self.assertTrue(layernorm_sp.runs_sp(ForwardMode.EXTEND))
|
|
# SP is prefill-only; decode must never engage it.
|
|
self.assertFalse(layernorm_sp.runs_sp(ForwardMode.DECODE))
|
|
# A disabled model never activates, even on EXTEND.
|
|
self.assertFalse(layernorm_sp.runs_sp(ForwardMode.EXTEND))
|
|
|
|
def test_runs_sp_ignores_the_active_flag(self):
|
|
"""Regression: the exit gather must not key off ``sp_active``.
|
|
|
|
``sp_active`` is written by Python inside the CUDA-graph-captured region,
|
|
so it is stale on graph replay. Callers outside that region (the exit
|
|
gather in LogitsProcessor) recompute with ``runs_sp`` instead; if that ever
|
|
starts consulting the flag, a replayed prefill skips the gather and feeds
|
|
sequence-sharded hidden states to the LM head.
|
|
"""
|
|
with get_flags().sp.override(enabled=True):
|
|
get_forward().set("sp_active", False)
|
|
self.assertTrue(layernorm_sp.runs_sp(ForwardMode.EXTEND))
|
|
reset_context()
|
|
|
|
|
|
class TestLayerNormSPValidation(CustomTestCase):
|
|
"""``validate_layernorm_sp`` is pure; pass config in directly."""
|
|
|
|
VALID = dict(
|
|
architecture="Qwen3ForCausalLM",
|
|
tp_size=2,
|
|
enable_dp_attention=False,
|
|
speculative_algorithm=None,
|
|
)
|
|
|
|
def test_valid_config_passes(self):
|
|
validate_layernorm_sp(**self.VALID) # must not raise
|
|
|
|
def test_rejects_unsupported_arch(self):
|
|
with self.assertRaisesRegex(ValueError, "only supported"):
|
|
validate_layernorm_sp(**{**self.VALID, "architecture": "LlamaForCausalLM"})
|
|
|
|
def test_rejects_tp_size_one(self):
|
|
with self.assertRaisesRegex(ValueError, "tp_size"):
|
|
validate_layernorm_sp(**{**self.VALID, "tp_size": 1})
|
|
|
|
def test_rejects_dp_attention(self):
|
|
with self.assertRaisesRegex(ValueError, "dp-attention"):
|
|
validate_layernorm_sp(**{**self.VALID, "enable_dp_attention": True})
|
|
|
|
def test_rejects_speculative(self):
|
|
with self.assertRaisesRegex(ValueError, "speculative"):
|
|
validate_layernorm_sp(**{**self.VALID, "speculative_algorithm": "EAGLE3"})
|
|
|
|
|
|
class TestLayerNormSPActiveFlag(CustomTestCase):
|
|
def tearDown(self):
|
|
reset_context()
|
|
|
|
def test_active_flag_is_registered_on_forward_flags(self):
|
|
"""``sp_active`` must stay a registered ForwardFlags slot.
|
|
|
|
``set()`` rejects names missing from ``ForwardFlags._DEFAULTS``, so this
|
|
fails if the slot is dropped, and it must also be in ``_GRAPH_VISIBLE``
|
|
because the participant linears read it under CUDA graph capture.
|
|
"""
|
|
from sglang.srt.runtime_context import ForwardFlags
|
|
|
|
self.assertIn("sp_active", ForwardFlags._DEFAULTS)
|
|
self.assertIn("sp_active", ForwardFlags._GRAPH_VISIBLE)
|
|
|
|
self.assertFalse(get_forward().sp_active) # default
|
|
get_forward().set("sp_active", True)
|
|
self.assertTrue(get_forward().sp_active)
|
|
get_forward().set("sp_active", False)
|
|
self.assertFalse(get_forward().sp_active)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|