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

[diffusion] refactor: refactor and simplify teacache for cachabledit and wanvideo (#16396)

Browse Source 
Co-authored-by: Brain97 <Brain97@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: blahblah <blahblah>
This commit is contained in:
blahblah
2026-01-21 15:42:45 +08:00
committed by GitHub
co-authored by Brain97 gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Mick blahblah <blahblah>
parent 0a9099e137
commit 0a7a2017a0
13 changed files with 702 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/test/server/perf_baselines.json
+67
View File
@@ -758,6 +758,73 @@
"expected_avg_denoise_ms": 104.84,
"expected_median_denoise_ms": 102.01
},
"wan2_1_t2v_1.3b_teacache_enabled": {
"stages_ms": {
"InputValidationStage": 0.07,
"TextEncodingStage": 2237.78,
"ConditioningStage": 0.01,
"TimestepPreparationStage": 2.1,
"LatentPreparationStage": 0.84,
"DenoisingStage": 13041.23,
"DecodingStage": 1274.63,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 879.71,
"1": 248.13,
"2": 246.48,
"3": 247.87,
"4": 249.38,
"5": 246.76,
"6": 250.42,
"7": 250.81,
"8": 250.98,
"9": 249.9,
"10": 246.72,
"11": 249.79,
"12": 250.46,
"13": 249.19,
"14": 247.55,
"15": 250.12,
"16": 247.57,
"17": 247.21,
"18": 247.32,
"19": 247.42,
"20": 248.21,
"21": 247.19,
"22": 247.72,
"23": 247.45,
"24": 247.9,
"25": 247.87,
"26": 247.18,
"27": 247.65,
"28": 246.91,
"29": 248.26,
"30": 247.82,
"31": 247.73,
"32": 247.38,
"33": 247.84,
"34": 247.46,
"35": 247.52,
"36": 247.94,
"37": 248.76,
"38": 248.01,
"39": 247.45,
"40": 247.84,
"41": 248.33,
"42": 247.41,
"43": 248.16,
"44": 248.18,
"45": 248.44,
"46": 248.65,
"47": 247.73,
"48": 247.48,
"49": 247.54
},
"expected_e2e_ms": 18382.19,
"expected_avg_denoise_ms": 260.76,
"expected_median_denoise_ms": 247.84
},
"wan2_1_t2v_1.3b": {
"stages_ms": {
"InputValidationStage": 0.07,
python/sglang/multimodal_gen/test/server/test_server_utils.py
+35 -1
View File
@@ -843,12 +843,18 @@ def get_generate_fn(
req_output_format = None # Not specified in current request
req_background = None # Not specified in current request
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
response = client.images.with_raw_response.generate(
model=model_path,
prompt=sampling_params.prompt,
n=n,
size=output_size,
response_format="b64_json",
extra_body=extra_body if extra_body else None,
)
result = response.parse()
validate_image(result.data[0].b64_json)
@@ -911,6 +917,11 @@ def get_generate_fn(
) # Not specified in current request
req_background = None # Not specified in current request
# Build extra_body for optional features
extra_body = {"num_frames": sampling_params.num_frames}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
images = [open(image_path, "rb") for image_path in image_paths]
try:
response = client.images.with_raw_response.edit(
@@ -921,7 +932,7 @@ def get_generate_fn(
size=output_size,
response_format="b64_json",
output_format=req_output_format,
extra_body={"num_frames": sampling_params.num_frames},
extra_body=extra_body,
)
finally:
for img in images:
@@ -1036,6 +1047,11 @@ def get_generate_fn(
if not sampling_params.prompt:
pytest.skip(f"{id}: no text prompt configured")
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
return _create_and_download_video(
client,
case_id,
@@ -1043,6 +1059,7 @@ def get_generate_fn(
prompt=sampling_params.prompt,
size=output_size,
seconds=video_seconds,
extra_body=extra_body if extra_body else None,
)
def generate_image_to_video(case_id, client) -> str:
@@ -1057,6 +1074,11 @@ def get_generate_fn(
if not image_path.exists():
pytest.skip(f"{id}: file missing: {image_path}")
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
with image_path.open("rb") as fh:
return _create_and_download_video(
client,
@@ -1066,11 +1088,18 @@ def get_generate_fn(
size=output_size,
seconds=video_seconds,
input_reference=fh,
extra_body=extra_body if extra_body else None,
)
def generate_text_url_image_to_video(case_id, client) -> str:
if not sampling_params.prompt or not sampling_params.image_path:
pytest.skip(f"{id}: no edit config")
# Build extra_body for optional features
extra_body = {"reference_url": sampling_params.image_path}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
return _create_and_download_video(
client,
case_id,
@@ -1097,6 +1126,11 @@ def get_generate_fn(
if not image_path.exists():
pytest.skip(f"{id}: file missing: {image_path}")
# Build extra_body for optional features
extra_body = {}
if sampling_params.enable_teacache:
extra_body["enable_teacache"] = True
with image_path.open("rb") as fh:
return _create_and_download_video(
client,
python/sglang/multimodal_gen/test/server/testcase_configs.py
+17
View File
@@ -191,6 +191,9 @@ class DiffusionSamplingParams:
num_outputs_per_prompt: int = 1
# TeaCache acceleration
enable_teacache: bool = False
@dataclass(frozen=True)
class DiffusionTestCase:
@@ -457,6 +460,20 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
prompt=T2V_PROMPT,
),
),
# TeaCache acceleration test for Wan video model
DiffusionTestCase(
"wan2_1_t2v_1.3b_teacache_enabled",
DiffusionServerArgs(
model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
modality="video",
warmup=0,
custom_validator="video",
),
DiffusionSamplingParams(
prompt=T2V_PROMPT,
enable_teacache=True,
),
),
# LoRA test case for single transformer + merge/unmerge API test
# Note: Uses dynamic_lora_path instead of lora_path to test LayerwiseOffload + set_lora interaction
# Server starts WITHOUT LoRA, then set_lora is called after startup (Wan models auto-enable layerwise offload)