minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] fix: preserve mapped courier tensor lifetime (#37965)

Browse Source 
Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
PengYuan
2026-09-05 07:41:47 +08:00
committed by GitHub
co-authored by Mick Qian
parent db89f639ef
commit 9c254df8cd
7 changed files with 148 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py
+6 -1
View File
@@ -1130,11 +1130,16 @@ class LayerwiseOffloadManager:
self._courier_inflight.discard(layer_idx)
self.prefetch_layer(layer_idx, non_blocking=False)
return
compute_stream = torch.get_device_module().current_stream()
compute_stream.wait_event(event)
with torch.inference_mode(False), torch.no_grad():
for name, gpu_tensor in tensors.items():
# wait_event orders the copy; record_stream keeps its storage
# live until the consuming kernels finish.
if gpu_tensor.device.type != "cpu":
gpu_tensor.record_stream(compute_stream)
target = self.get_target_with_name(name)
target.data = self._wrap_for_target(target, gpu_tensor)
torch.get_device_module().current_stream().wait_event(event)
self._courier_inflight.discard(layer_idx)
self._gpu_layers.add(layer_idx)
python/sglang/multimodal_gen/test/server/gpu_cases.py
+19 -22
View File
@@ -1074,6 +1074,20 @@ TWO_GPU_CASES = [
ring_degree=2,
),
),
DiffusionTestCase(
"qwen_image_t2i_2_gpus_extra_high",
DiffusionServerArgs(
model_path=DEFAULT_QWEN_IMAGE_MODEL_NAME_FOR_TEST,
# Cover the request-gated fused added-QKV path with the same ring setup.
ulysses_degree=1,
ring_degree=2,
),
replace(T2I_sampling_params, extras={"quality": "extra-high"}),
run_perf_check=False,
run_component_accuracy_check=False,
run_models_api_check=False,
run_t2v_input_reference_check=False,
),
DiffusionTestCase(
"zimage_image_t2i_2_gpus",
DiffusionServerArgs(
@@ -1132,20 +1146,6 @@ ONE_GPU_CASES += ONE_GPU_MODELOPT_FP8_CASES
TWO_GPU_CASES = _with_default_num_gpus(TWO_GPU_CASES, 2)
ONE_GPU_5090_PERF_CASE_IDS = frozenset(
{
"zimage_image_t2i",
"flux_2_klein_base_image_t2i",
"wan2_1_t2v_1.3b",
}
)
ONE_GPU_5090_SKIP_CONSISTENCY_CASE_IDS = frozenset(
{
"turbo_wan2_1_t2v_1.3b",
}
)
def _select_5090_canary_cases(case_ids: tuple[str, ...]) -> list[DiffusionTestCase]:
cases_by_id = {case.id: case for case in ONE_GPU_CASES}
missing = [case_id for case_id in case_ids if case_id not in cases_by_id]
@@ -1155,11 +1155,8 @@ def _select_5090_canary_cases(case_ids: tuple[str, ...]) -> list[DiffusionTestCa
return [
replace(
cases_by_id[case_id],
run_perf_check=case_id in ONE_GPU_5090_PERF_CASE_IDS,
run_consistency_check=(
cases_by_id[case_id].run_consistency_check
and case_id not in ONE_GPU_5090_SKIP_CONSISTENCY_CASE_IDS
),
run_perf_check=True,
run_consistency_check=True,
)
for case_id in case_ids
]
@@ -1187,8 +1184,8 @@ def _make_5090_flux_layerwise_cpu_offload_case() -> DiffusionTestCase:
output_size="512x512",
extras={"num_inference_steps": 4, "seed": 0},
),
run_perf_check=False,
run_consistency_check=False,
run_perf_check=True,
run_consistency_check=True,
run_component_accuracy_check=False,
run_models_api_check=False,
run_t2v_input_reference_check=False,
@@ -1256,7 +1253,7 @@ def _make_5090_h3_consumer_budget_case() -> DiffusionTestCase:
),
run_perf_check=True,
perf_repeat_requests=2,
run_consistency_check=False,
run_consistency_check=True,
run_component_accuracy_check=False,
run_models_api_check=False,
run_t2v_input_reference_check=False,
python/sglang/multimodal_gen/test/server/perf_baselines/5090.json
+73 -19
View File
@@ -2,8 +2,8 @@
"metadata": {
"model": "Diffusion Server",
"hardware": "CI RTX 5090 pool",
"description": "Reference numbers tightened from RTX 5090 multimodal_gen CI run 31935214180 at source revision 2ee0d38a8553d186faad9c50893480e3c1d79270.",
"last_updated": "2026-08-16"
"description": "Reference numbers from RTX 5090 multimodal_gen CI runs 31935214180 and 33854592181.",
"last_updated": "2026-09-04"
},
"tolerances": {
"long_term": {
@@ -88,6 +88,31 @@
"runtime_peak_vram_mb": 18246.0,
"estimated_full_test_time_s": 160.9
},
"turbo_wan2_1_t2v_1.3b": {
"stages_ms": {
"InputValidationStage": 0.11,
"TextEncodingStage": 634.42,
"TimestepPreparationStage": 65.02,
"LatentPreparationStage": 0.23,
"DmdDenoisingStage": 577.26,
"DecodingStage": 493.29,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 91.5,
"1": 160.17,
"2": 161.71,
"3": 161.1
},
"expected_e2e_ms": 2155.59,
"expected_avg_denoise_ms": 143.62,
"expected_median_denoise_ms": 160.64,
"load_peak_vram_mb": 928.0,
"runtime_peak_vram_mb": 14216.0,
"load_peak_host_anon_mb": 6759.33,
"runtime_peak_host_anon_mb": 6759.33,
"estimated_full_test_time_s": 200.3
},
"zimage_image_t2i": {
"stages_ms": {
"DecodingStage": 7.11,
@@ -112,6 +137,30 @@
"runtime_peak_vram_mb": 17776.0,
"estimated_full_test_time_s": 329.8
},
"flux_image_t2i_layerwise_cpu_offload_5090": {
"stages_ms": {
"InputValidationStage": 0.11,
"TextEncodingStage": 204.63,
"LatentPreparationStage": 0.72,
"TimestepPreparationStage": 71.6,
"DenoisingStage": 3603.59,
"DecodingStage": 37.95
},
"denoise_step_ms": {
"0": 400.73,
"1": 682.43,
"2": 876.57,
"3": 877.12
},
"expected_e2e_ms": 4604.15,
"expected_avg_denoise_ms": 709.21,
"expected_median_denoise_ms": 779.5,
"load_peak_vram_mb": 4916.0,
"runtime_peak_vram_mb": 19146.0,
"load_peak_host_anon_mb": 10559.66,
"runtime_peak_host_anon_mb": 10559.66,
"estimated_full_test_time_s": 172.8
},
"lingbot_video_moe_t2v": {
"stages_ms": {
"InputValidationStage": 0.1,
@@ -136,28 +185,33 @@
},
"minimax_h3_t2va_consumer_budget_1gpu_5090": {
"stages_ms": {
"MiniMaxH3TextEncodingStage": 11900.0,
"MiniMaxH3DenoisingStage": 76100.0,
"MiniMaxH3DecodingStage": 13000.0
"InputValidationStage": 0.06,
"MiniMaxH3PartitionAdmissionStage": 0.03,
"MiniMaxH3TextEncodingStage": 1825.86,
"MiniMaxH3VisualEncodingStage": 0.09,
"MiniMaxH3AudioEncodingStage": 0.06,
"MiniMaxH3LatentPreparationStage": 6.56,
"MiniMaxH3TimestepPreparationStage": 0.28,
"MiniMaxH3DenoisingStage": 34839.69,
"MiniMaxH3DecodingStage": 5257.87
},
"denoise_step_ms": {
"0": 14000.0,
"1": 14000.0,
"2": 14000.0,
"3": 14000.0,
"4": 14000.0,
"5": 14000.0,
"6": 14000.0,
"7": 14000.0
"0": 4941.69,
"1": 4863.02,
"2": 4223.75,
"3": 4384.88,
"4": 6182.67,
"5": 5068.86,
"6": 4579.34
},
"expected_e2e_ms": 105000.0,
"expected_avg_denoise_ms": 12000.0,
"expected_median_denoise_ms": 11000.0,
"load_peak_vram_mb": 6000.0,
"runtime_peak_vram_mb": 12288.0,
"expected_e2e_ms": 42763.29,
"expected_avg_denoise_ms": 4892.03,
"expected_median_denoise_ms": 4863.02,
"load_peak_vram_mb": 5052.0,
"runtime_peak_vram_mb": 12278.0,
"load_peak_host_anon_mb": 32768.0,
"runtime_peak_host_anon_mb": 32768.0,
"estimated_full_test_time_s": 1050.0
"estimated_full_test_time_s": 283.0
}
}
}
python/sglang/multimodal_gen/test/server/perf_baselines/h100.json
+8
View File
@@ -175,6 +175,14 @@
"runtime_peak_vram_mb": 46938.0,
"estimated_full_test_time_s": 133.2
},
"qwen_image_t2i_2_gpus_extra_high": {
"stages_ms": {},
"denoise_step_ms": {},
"expected_e2e_ms": 0.0,
"expected_avg_denoise_ms": 0.0,
"expected_median_denoise_ms": 0.0,
"estimated_full_test_time_s": 57.9
},
"ideogram4_fp8_t2i": {
"stages_ms": {
"InputValidationStage": 0.06,
python/sglang/multimodal_gen/test/test_utils.py
+1 -1
View File
@@ -40,7 +40,7 @@ logger = init_logger(__name__)
# NPU/ascend) is read from sgl-project/ci-data-diffusion, where the GT-gen workflows
# publish.
SGL_TEST_FILES_CI_DATA_REPO = "sgl-project/ci-data-diffusion"
SGL_TEST_FILES_CI_DATA_REVISION = "15b30030ef980756788ab40072f9223fe21a5526"
SGL_TEST_FILES_CI_DATA_REVISION = "11783b3fbd8ebb1e3509cc4590c1fff476e65511"
# The NPU pin is kept as a separate branch so ascend GT can be bumped independently
# when it's regenerated on its own cadence.
python/sglang/multimodal_gen/test/unit/test_layerwise_offload.py
+30
View File
@@ -1386,6 +1386,36 @@ def test_mapped_layers_ship_through_the_courier(tmp_path, monkeypatch):
)
def test_collected_mapped_weights_record_the_compute_stream(monkeypatch):
"""Courier tensors must stay live until their compute-stream kernels finish."""
compute_stream = _FakeStream()
recorded_streams = []
gpu_tensor = SimpleNamespace(
device=torch.device("cuda"), record_stream=recorded_streams.append
)
monkeypatch.setattr(
_FakeDeviceModule, "current_stream", staticmethod(lambda: compute_stream)
)
monkeypatch.setattr(
layerwise_offload_mod.torch, "get_device_module", lambda: _FakeDeviceModule
)
manager = object.__new__(LayerwiseOffloadManager)
target = torch.nn.Parameter(torch.zeros(1))
manager._mapped_courier = SimpleNamespace(
collect=lambda _layer_idx: (_FakeEvent(), {"weight": gpu_tensor})
)
manager._named_parameters = {"weight": target}
manager._named_buffers = {}
manager._wrap_for_target = lambda _target, _tensor: target.data
manager._courier_inflight = {0}
manager._gpu_layers = set()
manager._collect_mapped_layer(0)
assert recorded_streams == [compute_stream]
def test_the_courier_kill_switch_forces_the_synchronous_path(tmp_path, monkeypatch):
if not pathlib.Path("/proc/self/maps").exists():
pytest.skip("needs /proc to tell a mapping from anonymous memory")
python/sglang/multimodal_gen/test/unit/test_suite_partitioning.py
+11
View File
@@ -18,6 +18,7 @@ from sglang.multimodal_gen.test.run_suite import (
from sglang.multimodal_gen.test.server.gpu_cases import (
PARAMETRIZED_CASE_GROUPS,
STANDALONE_FILES,
TWO_GPU_CASES,
)
@@ -108,3 +109,13 @@ def test_failing_cases_do_not_skip_the_shards_standalone_files(monkeypatch, tmp_
assert executed_standalone == [standalone_rel]
assert exit_code == 1
def test_qwen_quality_variants_use_the_same_generation_request():
cases = {case.id: case for case in TWO_GPU_CASES}
lossless = cases["qwen_image_t2i_2_gpus"].sampling_params
extra_high = cases["qwen_image_t2i_2_gpus_extra_high"].sampling_params
assert extra_high.prompt == lossless.prompt
assert extra_high.output_size == lossless.output_size
assert extra_high.extras == {"quality": "extra-high"}