[Feature] Add AllGather-KV sequence-parallel attention str… (#4968)

Signed-off-by: zhuhh97 <zhuhonghan@huawei.com>
Signed-off-by: dph97 <nenbaying@163.com>
Signed-off-by: zhuhh97 <zhh@unknown>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: dengyunyang <584797741@qq.com>
This commit is contained in:
authored and GitHub committed 2026-07-24 14:45:26 +00:00
1 parent 8927ec057c
commit 84bbd57e4d
33 files changed
+1181 -135

No files matched your search

@@ -3,6 +3,7 @@ from __future__ import annotations
import base64
import copy
import io
import logging
import os
import tempfile
from pathlib import Path
@@ -18,6 +19,8 @@ from tests.helpers.runtime import OmniServer
pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
logger = logging.getLogger(__name__)
MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
SEED = 42
NUM_INFERENCE_STEPS = 50
@@ -53,11 +56,12 @@ _DEPLOY_CONFIG = {
"parallel_config": {
"pipeline_parallel_size": 1,
"data_parallel_size": 1,
"tensor_parallel_size": 4,
"tensor_parallel_size": 1,
"enable_expert_parallel": True,
"sequence_parallel_size": 1,
"ulysses_degree": 1,
"ring_degree": 1,
"allgather_degree": 4,
"cfg_parallel_size": 1,
"vae_patch_parallel_size": 1,
"use_hsdp": False,
@@ -83,9 +87,47 @@ def _devices() -> str:
def _write_deploy_config(path: Path) -> None:
config = copy.deepcopy(_DEPLOY_CONFIG)
devices = _devices()
num_devices = len(devices.split(","))
config["stages"][0]["devices"] = devices
config["stages"][0]["parallel_config"]["tensor_parallel_size"] = len(devices.split(","))
path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
parallel_config = config["stages"][0]["parallel_config"]
if parallel_config.get("allgather_degree", 1) > 1:
parallel_config["allgather_degree"] = num_devices
elif parallel_config.get("tensor_parallel_size", 1) > 1:
parallel_config["tensor_parallel_size"] = num_devices
sequence_parallel_size = (
parallel_config.get("allgather_degree", 1)
if parallel_config.get("allgather_degree", 1) > 1
else parallel_config.get("ulysses_degree", 1) * parallel_config.get("ring_degree", 1)
)
configured_world_size = (
parallel_config.get("pipeline_parallel_size", 1)
* parallel_config.get("data_parallel_size", 1)
* parallel_config.get("tensor_parallel_size", 1)
* sequence_parallel_size
* parallel_config.get("cfg_parallel_size", 1)
)
assert configured_world_size == num_devices, (
f"Configured parallel world size ({configured_world_size}) does not match "
f"the number of devices ({num_devices}): {parallel_config}"
)
config_text = yaml.dump(config, default_flow_style=False, sort_keys=False)
path.write_text(config_text)
logger.info(
"HunyuanImage3 pixel accuracy deploy config:\n"
" path=%s\n"
" CUDA_VISIBLE_DEVICES=%s\n"
" devices=%s\n"
" tensor_parallel_size=%s\n"
" allgather_degree=%s\n%s",
path,
os.environ.get("CUDA_VISIBLE_DEVICES", "<unset>"),
devices,
parallel_config.get("tensor_parallel_size", 1),
parallel_config.get("allgather_degree", 1),
config_text,
)
def _run_vllm_omni_hunyuan_image3_online(*, model: str, deploy_config: str, output_path: Path) -> Image.Image: