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

[NPU] [Diffusion] Fix critical Ascend NPU Diffusion regression/bugs & restore 2-NPU CI testcase (#34855)

Browse Source 
Co-authored-by: Elizaveta Martirosian <elizabet3000@mail.ru>
Co-authored-by: Arseniy Mironov <98156294+Napkin-AI@users.noreply.github.com>
Co-authored-by: Alexandr <117110413+Allor-maker@users.noreply.github.com>
Co-authored-by: P_Alex_Tr <aleksandr.smyshlaev@yandex.ru>
This commit is contained in:
Артем Савкин
2026-08-22 18:32:20 +03:00
committed by GitHub
co-authored by Elizaveta Martirosian Arseniy Mironov Alexandr P_Alex_Tr
parent 0b064e3739
commit b98d472158
20 changed files with 356 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/flash_attention.py
-1
View File
@@ -257,7 +257,6 @@ def flash_attn_varlen_func(
ver=3,
out=None,
):
if ver == 3:
return fa3_flash_attn_varlen_func(
q,
python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py
+205 -1
View File
@@ -1,3 +1,4 @@
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any
@@ -15,6 +16,143 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
def _packed_boundaries(
cu_seqlens: torch.Tensor,
cu_seqlens_host: Sequence[int] | None,
total_tokens: int,
name: str,
) -> tuple[int, ...]:
if cu_seqlens is None:
raise ValueError(f"{name} is required for NPU packed attention")
if cu_seqlens.ndim != 1 or cu_seqlens.dtype not in (
torch.int32,
torch.int64,
):
raise ValueError(f"{name} must be a 1D int32 or int64 tensor")
if cu_seqlens_host is not None and len(cu_seqlens_host) != cu_seqlens.numel():
raise ValueError(f"{name} and its host copy must have the same length")
boundaries = tuple(
int(value)
for value in (
cu_seqlens.tolist() if cu_seqlens_host is None else cu_seqlens_host
)
)
if len(boundaries) < 2 or boundaries[0] != 0:
raise ValueError(f"{name} must start with 0 and contain at least one sequence")
if boundaries[-1] != total_tokens:
raise ValueError(
f"{name} must end at the packed token count {total_tokens}, "
f"got {boundaries[-1]}"
)
if any(stop < start for start, stop in zip(boundaries[:-1], boundaries[1:])):
raise ValueError(f"{name} must be non-decreasing")
return boundaries
def fused_infer_attention_varlen(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens_q: torch.Tensor,
cu_seqlens_k: torch.Tensor,
*,
cu_seqlens_q_host: Sequence[int] | None = None,
cu_seqlens_k_host: Sequence[int] | None = None,
softmax_scale: float | None = None,
return_softmax_lse: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
tensors = {"q": q, "k": k, "v": v}
invalid_layouts = [name for name, tensor in tensors.items() if tensor.ndim != 3]
if invalid_layouts:
raise ValueError(
"NPU packed attention requires q, k, and v in [T, N, D] layout; "
f"invalid tensors: {', '.join(invalid_layouts)}"
)
invalid_devices = [
name
for name, tensor in tensors.items()
if tensor.device.type != "npu" or tensor.device != q.device
]
if invalid_devices:
raise ValueError(
"NPU packed attention requires q, k, and v on the same NPU; "
f"invalid tensors: {', '.join(invalid_devices)}"
)
if not (q.dtype == k.dtype == v.dtype):
raise ValueError(
"NPU packed attention requires q, k, and v with the same dtype"
)
if k.shape[:2] != v.shape[:2]:
raise ValueError(
"NPU packed attention requires matching K/V token and head counts"
)
if q.shape[-1] != k.shape[-1]:
raise ValueError("NPU packed attention requires matching Q/K head dimensions")
q_boundaries = _packed_boundaries(
cu_seqlens_q, cu_seqlens_q_host, q.shape[0], "cu_seqlens_q"
)
k_boundaries = _packed_boundaries(
cu_seqlens_k, cu_seqlens_k_host, k.shape[0], "cu_seqlens_k"
)
if len(q_boundaries) != len(k_boundaries):
raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same batch")
q_nonempty = [
stop > start for start, stop in zip(q_boundaries[:-1], q_boundaries[1:])
]
k_nonempty = [
stop > start for start, stop in zip(k_boundaries[:-1], k_boundaries[1:])
]
if q_nonempty != k_nonempty:
raise NotImplementedError(
"NPU packed attention does not support a sequence that is empty only "
"on the query or key/value side"
)
actual_seq_lengths = [
stop for stop, nonempty in zip(q_boundaries[1:], q_nonempty) if nonempty
]
actual_seq_lengths_kv = [
stop for stop, nonempty in zip(k_boundaries[1:], k_nonempty) if nonempty
]
if not actual_seq_lengths:
output = torch.empty_like(q)
if return_softmax_lse:
lse = torch.empty(
(q.shape[1], q.shape[0]), dtype=torch.float32, device=q.device
)
return output, lse
return output
if not (q.is_contiguous() and k.is_contiguous() and v.is_contiguous()):
if q.shape == k.shape == v.shape:
q, k, v = torch.stack((q, k, v), dim=0).unbind(0)
else:
q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
output, lse = torch.ops.npu.npu_fused_infer_attention_score(
q,
k,
v,
num_heads=q.shape[1],
num_key_value_heads=k.shape[1],
scale=q.shape[-1] ** -0.5 if softmax_scale is None else softmax_scale,
input_layout="TND",
actual_seq_lengths=actual_seq_lengths,
actual_seq_lengths_kv=actual_seq_lengths_kv,
softmax_lse_flag=return_softmax_lse,
)
if not return_softmax_lse:
return output
if lse.shape != (q.shape[0], q.shape[1], 1):
raise RuntimeError(
"Unexpected Ascend TND softmax LSE shape: "
f"expected {(q.shape[0], q.shape[1], 1)}, got {tuple(lse.shape)}"
)
return output, lse.squeeze(-1).transpose(0, 1).contiguous()
@dataclass
class AscendFAMetadata:
pass
@@ -106,5 +244,71 @@ class AscendFAImpl(AttentionImpl):
)
output = output.transpose(1, 2)
if return_softmax_lse:
return output, lse
return output, lse.squeeze(-1)
return output
def forward_varlen(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
*,
cu_seqlens: torch.Tensor,
max_seqlen: int,
cu_seqlens_host: tuple[int, ...] | None = None,
) -> torch.Tensor:
del max_seqlen
if self.causal:
bounds = (
cu_seqlens_host
if cu_seqlens_host is not None
else tuple(int(item) for item in cu_seqlens.tolist())
)
output = torch.empty_like(query)
for start, stop in zip(bounds[:-1], bounds[1:]):
if start == stop:
continue
segment = self.forward(
query[start:stop].unsqueeze(0),
key[start:stop].unsqueeze(0),
value[start:stop].unsqueeze(0),
None,
)
output[start:stop].copy_(segment[0])
return output
return fused_infer_attention_varlen(
query,
key,
value,
cu_seqlens,
cu_seqlens,
cu_seqlens_q_host=cu_seqlens_host,
cu_seqlens_k_host=cu_seqlens_host,
softmax_scale=self.softmax_scale,
)
def forward_ring_kv_chunk(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Run one Ascend TND ring chunk and return LSE in ``[H, Tq]``."""
cu_seqlens_q = torch.tensor(
[0, query.shape[0]], dtype=torch.int32, device=query.device
)
cu_seqlens_k = torch.tensor(
[0, key.shape[0]], dtype=torch.int32, device=key.device
)
return fused_infer_attention_varlen(
query,
key,
value,
cu_seqlens_q,
cu_seqlens_k,
cu_seqlens_q_host=(0, query.shape[0]),
cu_seqlens_k_host=(0, key.shape[0]),
softmax_scale=self.softmax_scale,
return_softmax_lse=True,
)
python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py
+15
View File
@@ -209,6 +209,21 @@ class AttentionImpl(ABC, Generic[T]):
f"{type(self).__name__} does not implement packed varlen attention"
)
def forward_ring_kv_chunk(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Attend local queries to one rotated KV chunk for ring merging.
Inputs use packed ``[T, H, D]`` layout. The returned attention output
has the query shape and softmax LSE uses ``[H, Tq]`` layout.
"""
raise NotImplementedError(
f"{type(self).__name__} does not implement ring KV-chunk attention"
)
def wrap_attention_impl_forward(attn_impl: AttentionImpl) -> AttentionImpl:
return wrap_method_with_debug_kernel_once(
python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
+34
View File
@@ -472,3 +472,37 @@ class FlashAttentionImpl(AttentionImpl):
ver=fa_ver,
)
return output[0] if isinstance(output, tuple) else output
def forward_ring_kv_chunk(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Run one non-causal FlashAttention ring chunk with LSE output."""
cu_seqlens_q = torch.tensor(
[0, query.shape[0]], dtype=torch.int32, device=query.device
)
cu_seqlens_k = torch.tensor(
[0, key.shape[0]], dtype=torch.int32, device=key.device
)
result = flash_attn_varlen_func(
query,
key,
value,
cu_seqlens_q=cu_seqlens_q,
cu_seqlens_k=cu_seqlens_k,
max_seqlen_q=query.shape[0],
max_seqlen_k=key.shape[0],
softmax_scale=self.softmax_scale,
causal=False,
ver=fa_ver,
return_softmax_lse=True,
)
if not isinstance(result, tuple):
raise RuntimeError(
"FlashAttention did not return the softmax LSE required by ring "
"attention"
)
output, softmax_lse, *_ = result
return output, softmax_lse
python/sglang/multimodal_gen/runtime/layers/attention/layer.py
+1 -1
View File
@@ -1303,7 +1303,7 @@ class USPAttention(nn.Module):
q.squeeze(0),
k.squeeze(0),
v.squeeze(0),
softmax_scale=self.softmax_scale,
attn_impl=self.attn_impl,
real_seq_len=int(attn_mask_meta["pad_start"]),
ring_ws=get_ring_parallel_world_size(),
)
python/sglang/multimodal_gen/runtime/layers/usp.py
+2 -23
View File
@@ -8,7 +8,6 @@ import torch.distributed as dist
import torch.distributed._functional_collectives as ft_c
from torch.distributed.tensor.experimental._attention import _cp_options
from sglang.kernels.ops.attention.flash_attention import flash_attn_varlen_func
from sglang.kernels.ops.diffusion import pack_qkv_destination_major, usp_merge_heads
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_ring_ctx,
@@ -16,9 +15,6 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_ulysses_parallel_rank,
get_ulysses_parallel_world_size,
)
from sglang.multimodal_gen.runtime.layers.attention.backends import (
flash_attn as _fa_backend,
)
from sglang.srt.utils.common import torch_release
_cp_options.enable_load_balance = False
@@ -821,7 +817,7 @@ def _ring_attention_varlen(
k: torch.Tensor,
v: torch.Tensor,
*,
softmax_scale: float,
attn_impl: "AttentionImpl",
real_seq_len: int,
ring_ws: int,
) -> torch.Tensor:
@@ -860,7 +856,6 @@ def _ring_attention_varlen(
kv_bufs = [kv0, torch.empty_like(kv0)]
cur = 0
q_cu = torch.tensor([0, ring_chunk_len], dtype=torch.int32, device=q.device)
out_acc: torch.Tensor | None = None
lse_acc: torch.Tensor | None = None
pending_ops = None
@@ -891,27 +886,11 @@ def _ring_attention_varlen(
max(real_seq_len - src_rank * ring_chunk_len, 0), ring_chunk_len
)
if remote_used > 0:
k_cu = torch.tensor([0, remote_used], dtype=torch.int32, device=q.device)
result = flash_attn_varlen_func(
step_out, step_lse = attn_impl.forward_ring_kv_chunk(
q,
kv_bufs[cur][0, :remote_used],
kv_bufs[cur][1, :remote_used],
cu_seqlens_q=q_cu,
cu_seqlens_k=k_cu,
max_seqlen_q=ring_chunk_len,
max_seqlen_k=remote_used,
softmax_scale=softmax_scale,
causal=False,
ver=_fa_backend.fa_ver,
return_softmax_lse=True,
)
if not isinstance(result, tuple):
raise RuntimeError(
"flash_attn_varlen_func did not return softmax_lse; ring "
"parallelism requires a backend that supports "
"return_softmax_lse=True."
)
step_out, step_lse, *_ = result
out_acc, lse_acc = _ring_merge_attention(
out_acc, lse_acc, step_out, step_lse
)
python/sglang/multimodal_gen/runtime/managers/memory_managers/host_memory_budget.py
+7 -3
View File
@@ -239,12 +239,16 @@ def module_weight_bytes(module) -> int:
seen: set[int] = set()
total = 0
for tensor in list(module.parameters()) + list(module.buffers()):
storage = tensor.untyped_storage()
pointer = storage.data_ptr()
try:
storage = tensor.untyped_storage()
pointer = storage.data_ptr()
storage_bytes = storage.nbytes()
except RuntimeError:
continue
if pointer == 0 or pointer in seen:
continue
seen.add(pointer)
total += storage.nbytes()
total += storage_bytes
return total
python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+1 -1
View File
@@ -552,7 +552,7 @@ def _minimax_h3_attention_core_impl(
q,
k,
v,
softmax_scale=attention.softmax_scale,
attn_impl=attention._attention_impl,
real_seq_len=max_seqlen,
ring_ws=ring_ws,
)
python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py
+1 -1
View File
@@ -179,7 +179,7 @@ class SelfAttention(nn.Module):
# USPAttention handles SP communication internally; the tail meta keeps
# SP padding out of the softmax.
out = self.attn(q, k, v, attn_mask_meta=attn_mask_meta)
out = out.view(b, s, -1)
out = out.reshape(b, s, -1)
out, _ = self.o(out)
return out
python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json
+14 -9
View File
@@ -73,7 +73,7 @@
"flux_2_image_t2i_2npu": {
"stages_ms": {
"InputValidationStage": 0.08,
"TextEncodingStage": 192.4,
"TextEncodingStage": 363.09,
"ImageVAEEncodingStage": 0.01,
"LatentPreparationStage": 0.97,
"TimestepPreparationStage": 34.65,
@@ -134,7 +134,8 @@
},
"expected_e2e_ms": 46557.7,
"expected_avg_denoise_ms": 872.7,
"expected_median_denoise_ms": 905.81
"expected_median_denoise_ms": 905.81,
"estimated_full_test_time_s": 488.4
},
"wan2_1_t2v_1.3b_1_npu": {
"stages_ms": {
@@ -257,7 +258,8 @@
},
"expected_e2e_ms": 193947.19,
"expected_avg_denoise_ms": 4691.07,
"expected_median_denoise_ms": 4773.22
"expected_median_denoise_ms": 4773.22,
"estimated_full_test_time_s": 987.8
},
"qwen_image_t2i_2npu": {
"stages_ms": {
@@ -322,14 +324,15 @@
},
"expected_e2e_ms": 34362.34,
"expected_avg_denoise_ms": 610.41,
"expected_median_denoise_ms": 615.39
"expected_median_denoise_ms": 615.39,
"estimated_full_test_time_s": 275.2
},
"ernie_image_t2i_1npu": {
"stages_ms": {
"InputValidationStage": 0.07,
"PromptEnhancementStage": 8787.2,
"TextEncodingStage": 35.87,
"TimestepPreparationStage": 32.74,
"TimestepPreparationStage": 246.30,
"LatentPreparationStage": 0.2,
"DenoisingStage": 47471.93,
"DecodingStage": 42.24
@@ -572,7 +575,8 @@
},
"expected_e2e_ms": 109909.47,
"expected_avg_denoise_ms": 2636.6,
"expected_median_denoise_ms": 2705.05
"expected_median_denoise_ms": 2705.05,
"estimated_full_test_time_s": 239.7
},
"mova_360p_ti2va_2npu": {
"stages_ms": {
@@ -645,13 +649,13 @@
"stages_ms": {
"InputValidationStage": 3.8,
"TextEncodingStage": 1029.19,
"LTX2TextConnectorStage": 56.3,
"LTX2TextConnectorStage": 654.62,
"LTX2SigmaPreparationStage": 0.19,
"TimestepPreparationStage": 33.52,
"LTX2AVLatentPreparationStage": 0.45,
"LTX2ImageEncodingStage": 93.32,
"LTX2AVDenoisingStage": 29672.06,
"LTX2AVDecodingStage": 560.9
"LTX2AVDecodingStage": 1240.88
},
"denoise_step_ms": {
"0": 544.81,
@@ -697,7 +701,8 @@
},
"expected_e2e_ms": 31484.37,
"expected_avg_denoise_ms": 741.58,
"expected_median_denoise_ms": 746.28
"expected_median_denoise_ms": 746.28,
"estimated_full_test_time_s": 255.2
}
}
}
python/sglang/multimodal_gen/test/test_utils.py
+14 -5
View File
@@ -45,7 +45,7 @@ SGL_TEST_FILES_CI_DATA_REVISION = "15b30030ef980756788ab40072f9223fe21a5526"
# The NPU pin is kept as a separate branch so ascend GT can be bumped independently
# when it's regenerated on its own cadence.
if current_platform.is_npu():
SGL_TEST_FILES_CI_DATA_REVISION = "d180ad38872dff3d1ad03e4610cffcda874d3eb8"
SGL_TEST_FILES_CI_DATA_REVISION = "8e3d717e65fb87339c2974382a092a731669f884"
SGL_TEST_FILES_CONSISTENCY_GT_ROOT = (
"https://raw.githubusercontent.com/"
@@ -1632,20 +1632,29 @@ def _remote_file_exists(url: str) -> bool | None:
def _load_remote_gt_image(url: str) -> np.ndarray:
last_error: Exception | None = None
for _ in range(3):
attempts = 3
for attempt in range(1, attempts + 1):
try:
resp = requests.get(url, timeout=60)
try:
if resp.status_code == 200:
image = Image.open(io.BytesIO(resp.content)).convert("RGB")
return np.array(image)
with Image.open(io.BytesIO(resp.content)) as image:
return np.array(image.convert("RGB"))
last_error = FileNotFoundError(f"GT image not found: {url}")
if resp.status_code not in (403, 429) and resp.status_code < 500:
break
finally:
resp.close()
except requests.RequestException as exc:
except (OSError, ValueError, requests.RequestException) as exc:
last_error = exc
if attempt < attempts:
logger.warning(
"GT image download failed (attempt %d/%d), retrying: %s",
attempt,
attempts,
url,
)
time.sleep(attempt)
raise FileNotFoundError(f"GT image not found: {url}") from last_error
python/sglang/multimodal_gen/test/unit/test_host_memory_budget.py
+20
View File
@@ -239,6 +239,26 @@ class TestModuleWeightBytes:
module.register_buffer("b", backing[512:])
assert module_weight_bytes(module) == 4096
def test_invalid_storage_is_skipped(self):
"""Offloaded models may expose a tensor whose storage cannot be queried."""
class InvalidStorage:
def data_ptr(self):
raise RuntimeError("invalid python storage")
class InvalidStorageTensor:
def untyped_storage(self):
return InvalidStorage()
class ModuleWithInvalidStorage:
def parameters(self):
return iter((torch.ones(4), InvalidStorageTensor()))
def buffers(self):
return iter(())
assert module_weight_bytes(ModuleWithInvalidStorage()) == 4 * 4
class TestPinBenefit:
def test_a_stepped_component_counts_every_step(self):
python/sglang/multimodal_gen/test/unit/test_usp_ring_tail_pad.py
+4 -2
View File
@@ -22,6 +22,7 @@ class TestRingTailPadDispatch(unittest.TestCase):
obj.backend = AttentionBackendEnum.FA
obj.causal = False
obj.dropout_p = 0.0
obj.attn_impl = object()
return obj
def test_tail_pad_meta_reaches_the_ring_kernel(self):
@@ -30,12 +31,12 @@ class TestRingTailPadDispatch(unittest.TestCase):
meta = {"pad_start": 13, "pad_end": 16, "local_pad": 3}
seen = {}
def fake_ring(qc, kc, vc, *, softmax_scale, real_seq_len, ring_ws):
def fake_ring(qc, kc, vc, *, attn_impl, real_seq_len, ring_ws):
seen.update(
shape=tuple(qc.shape),
real=real_seq_len,
ws=ring_ws,
scale=softmax_scale,
impl=attn_impl,
)
return torch.ones_like(qc)
@@ -61,6 +62,7 @@ class TestRingTailPadDispatch(unittest.TestCase):
self.assertEqual(out.shape, q.shape)
self.assertEqual(seen["real"], 13)
self.assertEqual(seen["ws"], 2)
self.assertIs(seen["impl"], obj.attn_impl)
self.assertEqual(seen["shape"], (4, 2, 8))
# Last ring rank holds global rows [12, 16): row 13 onward is pad.
self.assertTrue(torch.all(out[0, 1:] == 0))
python/sglang/srt/hardware_backend/npu/utils.py
+7 -4
View File
@@ -101,12 +101,15 @@ def init_npu_backend():
logger.warning("NPU custom kernel packages unavailable: %s", e)
import torch_npu
from torch_npu.contrib import transfer_to_npu # noqa: F401
# Re-mock torch.cuda.is_available cuz transfer_to_npu mocks it True
torch.cuda.is_available = lambda: False
# These imports lead to unpredictable behavior in diffusion models
# and a significant reduction in performance.
if "sglang.multimodal_gen" not in sys.modules:
from torch_npu.contrib import transfer_to_npu # noqa: F401
torch_npu.npu.config.allow_internal_format = True
# Re-mock torch.cuda.is_available cuz transfer_to_npu mocks it True
torch.cuda.is_available = lambda: False
torch_npu.npu.config.allow_internal_format = True
torch_npu.npu.set_compile_mode(jit_compile=False)