From 740f45abcd2e6754d1cdded01ee9901d595ccac5 Mon Sep 17 00:00:00 2001 From: ooooooye <53851298+brandneway@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:47:03 +0800 Subject: [PATCH] [Perf][Diffusion] MiniMax-H3: opt into packed varlen attention on NPU to eliminate quadratic mask materialization (#5891) Signed-off-by: brandneway --- recipes/MiniMaxAI/MiniMax-H3-NPU.md | 71 ++++- tests/diffusion/attention/test_flash_attn.py | 262 ++++++++++++++++++ .../minimax_h3/test_minimax_h3_contract.py | 50 ++++ .../diffusion/attention/backends/abstract.py | 22 ++ .../attention/backends/flash_attn.py | 198 +++++++++++++ .../minimax_h3/minimax_h3_transformer.py | 31 ++- vllm_omni/platforms/npu/platform.py | 12 + 7 files changed, 638 insertions(+), 8 deletions(-) diff --git a/recipes/MiniMaxAI/MiniMax-H3-NPU.md b/recipes/MiniMaxAI/MiniMax-H3-NPU.md index 930411f1a..ffe8eee09 100755 --- a/recipes/MiniMaxAI/MiniMax-H3-NPU.md +++ b/recipes/MiniMaxAI/MiniMax-H3-NPU.md @@ -79,7 +79,8 @@ sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`. ### Multi-NPU: 768P combined-service configuration Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism -degree 8, native tiled VAE patch parallelism degree 8, and layerwise offload: +degree 8, native tiled VAE patch parallelism degree 8, and distributed +layerwise offload: ```bash export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 @@ -98,7 +99,7 @@ vllm serve "${MODEL}" \ --usp 8 \ --ring 1 \ --text-encoder-tp-size 8 \ - --enable-layerwise-offload \ + --enable-distributed-layerwise-offload \ --vae-parallel-mode tile \ --vae-use-tiling \ --vae-patch-parallel-size 8 \ @@ -112,6 +113,11 @@ H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no partition restart is required. Layerwise offload applies to both DiTs. +On Atlas 800I A3 (64 GB HBM per device) the combined service does not fit at +768P without offloading or sharding: use distributed layerwise offload (as +above) or HSDP — see +[§ Memory and attention optimizations](#memory-and-attention-optimizations-a3). + ### Optional optimizations Two independent optimizations may be enabled on top of the configuration @@ -134,6 +140,63 @@ Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks over the whole sequence, so ring parallelism would split away the keys it needs. Scale with `--usp` instead. +## Memory and attention optimizations (A3) + +### Fitting 768P into 64 GB HBM: distributed layerwise offload or HSDP + +Each Atlas 800I A3 NPU has 64 GB of HBM, which is not enough for the +combined MiniMax-H3 service at 768P. Enable **one** of the following at +server startup: + +**Distributed layerwise offload** — DiT layers are offloaded with parameters +gathered across the parallel group instead of replicated per rank: + +```bash + --enable-distributed-layerwise-offload +``` + +Measured peak NPU memory is about 45 GB per device for Ref2VA with a 13.88 s +input video generating a 15 s 768P (1344x768) output video. Host (CPU) memory +usage is high with this option. + +**HSDP** — hybrid sharded data parallelism for the DiT parameters. The +multi-stream memory-reuse knob is mandatory with this flag: + +```bash +export MULTI_STREAM_MEMORY_REUSE=2 +``` + +```bash + --use-hsdp +``` + +Host memory usage is small, but large shapes may still OOM; HBM usage +optimizations for this path are ongoing. + +### FLASH_ATTN backend with MindIE-SD + +Keep `--diffusion-attention-backend FLASH_ATTN` and install MindIE-SD (see +[§ Environment](#environment)). On NPU this backend carries most of the +memory-reduction and performance work: a mask-free packed varlen path that +never materializes the quadratic `full_qk` padding mask, and K/V prefix +slicing, both driven by the packed `cu_seqlens` metadata emitted by the H3 +transformer. + +### Optional: LaserAttention fused kernel + +For an additional attention speedup, select the Ascend LaserAttention fused +kernel before starting the server: + +```bash +export MINDIE_SD_FA_TYPE="ascend_laser_attention" +``` + +This requires the FLASH_ATTN backend and MindIE-SD. H3 automatically applies +exact power-of-two input pre-scaling (`laser_input_scale=256`) so the +kernel's fp16 workspace cannot overflow on outlier activations. Measured +speedup numbers will be added here. + + ## HTTP API examples The request format is identical to the GPU recipe; see @@ -154,8 +217,8 @@ configuration above: | Workload | Configuration | |----------|---------------| -| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile | -| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile | +| T2VA, 209 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile | +| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, distributed layerwise offload, Ulysses 8, VPP8 tile, regional compile | These measurements describe the validated shapes rather than a general throughput guarantee. diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index 53673dd58..e3f5355bf 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -9,6 +9,10 @@ This script tests two main scenarios: 2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding """ +import sys +from types import SimpleNamespace +from unittest.mock import Mock + import pytest import torch @@ -371,6 +375,264 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch): impl.forward_cuda(query, query, query, metadata) +# --------------------------------------------------------------------------- +# NPU mask-free packed attention (PR #5891). +# +# These tests cover the backend-side NPU branches that ``forward_fa_npu`` adds: +# - boundary resolution in ``_resolve_packed_seq_npu`` (pure Python, no device +# sync, no MindIE-SD import); +# - env dispatch in ``forward_fa_npu`` (which op a packed forward takes); +# - the laser input pre-scaling math in ``_forward_prefix_kv_slice_npu``. +# +# They run on CPU without a real NPU by injecting a fake ``mindiesd`` module +# into ``sys.modules`` (same pattern as ``test_paged_attention`` / +# ``benchmarks/conftest``). No production code is exercised for real kernels. + + +def _npu_impl(causal: bool = False) -> FlashAttentionImpl: + return FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=causal) + + +def _fake_mindiesd(monkeypatch, *, attention_forward=None, attention_forward_varlen=None): + """Install a stub ``mindiesd`` so local ``from mindiesd import ...`` works. + + Restored automatically by monkeypatch. Defaults are no-op callables so the + top-level import in ``forward_fa_npu`` succeeds even when the test mocks the + instance methods that would otherwise call them. + """ + fake = SimpleNamespace( + attention_forward=attention_forward or Mock(return_value=torch.zeros(1)), + attention_forward_varlen=attention_forward_varlen or Mock(return_value=torch.zeros(1)), + ) + monkeypatch.setitem(sys.modules, "mindiesd", fake) + return fake + + +# --- Test group A: boundary resolution (_resolve_packed_seq_npu) ------------- + + +def test_resolve_packed_seq_accepts_real_plus_pad_contract(): + impl = _npu_impl() + q = torch.randn(1, 8, 2, 4) + cu = torch.tensor([0, 5, 8], dtype=torch.int32) + extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5} + + assert impl._resolve_packed_seq_npu(q, q, extra) == ([5, 8], [5, 8]) + + +def test_resolve_packed_seq_single_document_no_padding(): + impl = _npu_impl() + q = torch.randn(1, 8, 2, 4) + cu = torch.tensor([0, 8], dtype=torch.int32) + extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 8, "max_seqlen_k": 8} + + assert impl._resolve_packed_seq_npu(q, q, extra) == ([8], [8]) + + +@pytest.mark.parametrize( + "case", + [ + "causal", + "missing_key", + "batch_ne_1", + "cu_too_long", + "cu_qk_mismatch", + "max_seqlen_not_int", + "used_gt_total", + "used_zero", + "pad_longer_than_real", + ], +) +def test_resolve_packed_seq_rejects_malformed_metadata(case): + """Any metadata outside the exact [real, pad] contract returns None so the + caller falls back to the masked path.""" + cu3 = torch.tensor([0, 5, 8], dtype=torch.int32) + cu2 = torch.tensor([0, 8], dtype=torch.int32) + cu4 = torch.tensor([0, 3, 5, 8], dtype=torch.int32) + cu_short = torch.tensor([0, 3, 8], dtype=torch.int32) + base = {"cu_seqlens_q": cu3, "cu_seqlens_k": cu3, "max_seqlen_q": 5, "max_seqlen_k": 5} + q = torch.randn(1, 8, 2, 4) + + impl = _npu_impl(causal=(case == "causal")) + if case == "missing_key": + extra = {"cu_seqlens_q": cu3, "max_seqlen_q": 5} # 3 of 4 keys absent + elif case == "batch_ne_1": + q = torch.randn(2, 8, 2, 4) + extra = base + elif case == "cu_too_long": + extra = {**base, "cu_seqlens_q": cu4, "cu_seqlens_k": cu4} + elif case == "cu_qk_mismatch": + extra = {**base, "cu_seqlens_q": cu3, "cu_seqlens_k": cu2} + elif case == "max_seqlen_not_int": + extra = {**base, "max_seqlen_q": 5.0} + elif case == "used_gt_total": + extra = {**base, "max_seqlen_q": 9} # 9 > total length 8 + elif case == "used_zero": + extra = {**base, "max_seqlen_q": 0, "max_seqlen_k": 0} + elif case == "pad_longer_than_real": + # real=3, pad=5: the real document must be the longer one. + extra = { + "cu_seqlens_q": cu_short, + "cu_seqlens_k": cu_short, + "max_seqlen_q": 3, + "max_seqlen_k": 3, + } + else: + extra = base + + assert impl._resolve_packed_seq_npu(q, q, extra) is None + + +# --- Test group B: env dispatch in forward_fa_npu (current behavior) -------- + + +def _npu_impl_with_mocked_paths() -> tuple[FlashAttentionImpl, torch.Tensor]: + """FlashAttentionImpl whose two packed paths are Mocks returning a sentinel. + + Lets a test assert which branch ``forward_fa_npu`` dispatched to without + running any real kernel. + """ + impl = _npu_impl() + sentinel = torch.tensor([1.0]) + impl._forward_varlen_packed_npu = Mock(return_value=sentinel) + impl._forward_prefix_kv_slice_npu = Mock(return_value=sentinel) + return impl, sentinel + + +def test_npu_env_unset_uses_varlen(monkeypatch): + monkeypatch.delenv("MINDIE_SD_FA_TYPE", raising=False) + _fake_mindiesd(monkeypatch) + impl, _ = _npu_impl_with_mocked_paths() + metadata = AttentionMetadata(extra={"npu_attn_varlen": True}) + + impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata) + + impl._forward_varlen_packed_npu.assert_called_once() + impl._forward_prefix_kv_slice_npu.assert_not_called() + + +def test_npu_env_laser_uses_slice(monkeypatch): + monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention") + _fake_mindiesd(monkeypatch) + impl, _ = _npu_impl_with_mocked_paths() + metadata = AttentionMetadata(extra={"npu_attn_varlen": True}) + + impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata) + + impl._forward_prefix_kv_slice_npu.assert_called_once() + impl._forward_varlen_packed_npu.assert_not_called() + + +@pytest.mark.parametrize("fa_type", ["prompt_flash_attn", "fused_attn_score", "ascend_laser_attenton", "garbage"]) +def test_npu_env_non_laser_value_falls_back_to_varlen(monkeypatch, fa_type): + """Locks in the CURRENT dispatch behavior (PR #5891). + + Any value other than ``ascend_laser_attention`` — including the valid + ``prompt_flash_attn``/``fused_attn_score`` and typos — is silently routed to + the TND varlen op, which cannot honor ``MINDIE_SD_FA_TYPE``. This is review + comment #1 on the PR and is intentionally captured here as a regression + net: update this test when the dispatch starts honoring/rejecting these. + """ + monkeypatch.setenv("MINDIE_SD_FA_TYPE", fa_type) + _fake_mindiesd(monkeypatch) + impl, _ = _npu_impl_with_mocked_paths() + metadata = AttentionMetadata(extra={"npu_attn_varlen": True}) + + impl.forward_fa_npu(torch.randn(1, 8, 2, 4), *[torch.randn(1, 8, 2, 4)] * 2, metadata) + + impl._forward_varlen_packed_npu.assert_called_once() + impl._forward_prefix_kv_slice_npu.assert_not_called() + + +def test_npu_varlen_opt_in_unset_takes_mask_path(monkeypatch): + """Models that do not set ``npu_attn_varlen`` (e.g. Wan/Cosmos) keep using + the existing masked attention_forward path; the new branches are skipped.""" + fake_forward = Mock(return_value=torch.zeros(1, 8, 2, 4)) + _fake_mindiesd(monkeypatch, attention_forward=fake_forward) + impl, _ = _npu_impl_with_mocked_paths() + metadata = AttentionMetadata( + attn_mask=torch.tensor([[True, True, True, True, True, False, False, False]]), + extra={}, # no npu_attn_varlen opt-in + ) + q = torch.randn(1, 8, 2, 4) + + out = impl.forward_fa_npu(q, q, q, metadata) + + impl._forward_varlen_packed_npu.assert_not_called() + impl._forward_prefix_kv_slice_npu.assert_not_called() + fake_forward.assert_called_once() + assert out is fake_forward.return_value + + +# --- Test group C: laser input pre-scaling in _forward_prefix_kv_slice_npu -- + + +def test_prefix_kv_slice_applies_laser_input_scaling(monkeypatch): + monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention") + captured: dict = {} + + def fake_attention_forward(query, key, value, **kwargs): + captured["query"] = query + captured["key"] = key + captured["value"] = value + captured.update(kwargs) + return torch.ones_like(query) # deterministic stand-in kernel output + + _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward) + impl = _npu_impl() # softmax_scale == 0.5 + q = torch.randn(1, 8, 2, 4) + cu = torch.tensor([0, 5, 8], dtype=torch.int32) + extra = { + "cu_seqlens_q": cu, + "cu_seqlens_k": cu, + "max_seqlen_q": 5, + "max_seqlen_k": 5, + "laser_input_scale": 256.0, + } + + out = impl._forward_prefix_kv_slice_npu(q, q, q, extra) + + # K/V sliced to the valid prefix (used_k == 5); q keeps full length. + assert captured["key"].shape == (1, 5, 2, 4) + assert captured["value"].shape == (1, 5, 2, 4) + assert captured["query"].shape == (1, 8, 2, 4) + assert captured["attn_mask"] is None + # Pre-divide q/k/v by the factor and compensate the kernel scale by its + # square (exact for power-of-two); output scaled back by the factor. + torch.testing.assert_close(captured["query"], q / 256.0) + torch.testing.assert_close(captured["key"], q[:, :5] / 256.0) + assert captured["scale"] == pytest.approx(0.5 * 256.0**2) + torch.testing.assert_close(out, torch.ones(1, 8, 2, 4) * 256.0) + + +def test_prefix_kv_slice_no_scaling_without_factor(monkeypatch): + monkeypatch.setenv("MINDIE_SD_FA_TYPE", "ascend_laser_attention") + captured: dict = {} + + def fake_attention_forward(query, key, value, **kwargs): + captured["query"] = query + captured.update(kwargs) + return torch.ones_like(query) + + _fake_mindiesd(monkeypatch, attention_forward=fake_attention_forward) + impl = _npu_impl() + q = torch.randn(1, 8, 2, 4) + cu = torch.tensor([0, 5, 8], dtype=torch.int32) + extra = { + "cu_seqlens_q": cu, + "cu_seqlens_k": cu, + "max_seqlen_q": 5, + "max_seqlen_k": 5, + # no laser_input_scale → no pre-scaling + } + + out = impl._forward_prefix_kv_slice_npu(q, q, q, extra) + + torch.testing.assert_close(captured["query"], q) + assert captured["scale"] == pytest.approx(0.5) + torch.testing.assert_close(out, torch.ones(1, 8, 2, 4)) + + if __name__ == "__main__": print("Running FlashAttention Padding Tests...") print("=" * 60) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 8cb4a5311..66fb52a54 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -492,6 +492,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask(): class FakeBackend: supports_prefix_kv_slicing = True + @classmethod + def supports_packed_mask_free(cls) -> bool: + return False + class FakeAttention(torch.nn.Module): attn_backend = FakeBackend @@ -526,6 +530,48 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask(): assert attention.attention.metadata.extra["valid_kv_length"] == 5 +def test_packed_attention_skips_mask_for_packed_mask_free_backend(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + ) + + class FakeBackend: + supports_prefix_kv_slicing = False + + @classmethod + def supports_packed_mask_free(cls) -> bool: + return True + + class FakeAttention(torch.nn.Module): + attn_backend = FakeBackend + + def __init__(self): + super().__init__() + self.metadata = None + + def forward(self, query, key, value, metadata): + self.metadata = metadata + return query + + attention = object.__new__(MiniMaxH3Attention) + torch.nn.Module.__init__(attention) + attention.attention = FakeAttention() + q = torch.randn(8, 2, 4) + + attention._run_packed_attention( + q, + q, + q, + cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32), + max_seqlen=5, + packed_total=8, + ) + + assert attention.attention.metadata.attn_mask is None + assert attention.attention.metadata.extra["valid_kv_length"] == 5 + assert attention.attention.metadata.extra["npu_attn_varlen"] is True + + def test_packed_attention_keeps_padding_mask_for_other_backends(): from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( MiniMaxH3Attention, @@ -534,6 +580,10 @@ def test_packed_attention_keeps_padding_mask_for_other_backends(): class FakeBackend: supports_prefix_kv_slicing = False + @classmethod + def supports_packed_mask_free(cls) -> bool: + return False + class FakeAttention(torch.nn.Module): attn_backend = FakeBackend diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 68e152063..f3a5f99ef 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -20,6 +20,17 @@ class AttentionBackend(ABC): # avoid a slower masked-attention plan when tail padding is not semantic. supports_prefix_kv_slicing: bool = False + @classmethod + def supports_packed_mask_free(cls) -> bool: + """Whether packed attention never reads attn_mask on this platform. + + When True, models that pack a [real, pad] two-document layout and + carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip + constructing the padding mask entirely. Backends whose mask-free + behavior is platform-dependent must check current_omni_platform. + """ + return False + # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted. # Platform resolution rejects an explicit selection outside this set, so a # hardware-specific backend fails before the model is built. @@ -121,6 +132,17 @@ class AttentionMetadata: # the packed cu_seqlens tensors. # "valid_kv_length": int — contiguous valid K/V prefix length for a # backend that advertises supports_prefix_kv_slicing. + # "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path + # (TND npu_fusion_attention driven by cu_seqlens, mask never read). + # Requires the [real, pad] two-document packing contract; see + # FlashAttentionImpl._forward_varlen_packed_npu. + # "laser_input_scale": float — model opt-in input pre-scale for the NPU + # ascend_laser_attention path. The kernel stores unscaled QK^T in an + # fp16 workspace, so outlier activations overflow 65504 into NaN rows; + # with this set (>1), q/k/v are divided by the factor before the op, + # the kernel scale_value is multiplied by its square, and the output + # is scaled back (exact for power-of-two factors). Absent means no + # pre-scaling. See FlashAttentionImpl._forward_prefix_kv_slice_npu. # Piecewise attention metadata (mixed causal/full masks). # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention. diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index facb1ace5..bdef1dd65 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os from functools import partial import torch @@ -16,6 +17,7 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( piecewise_attn, ) from vllm_omni.diffusion.config import get_current_diffusion_config_or_none +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -24,6 +26,15 @@ class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True supports_piecewise_spans: bool = True + @classmethod + def supports_packed_mask_free(cls) -> bool: + # CUDA: forward_cuda dispatches the packed cu_seqlens varlen path + # before ever reading attn_mask. NPU: forward_fa_npu honors the + # npu_attn_varlen opt-in and rebuilds the padding mask itself if the + # packed contract fails. XPU reads attn_mask, so models must keep + # constructing it there. + return current_omni_platform.is_cuda() or current_omni_platform.is_npu() + @classmethod def supports_attention_mask(cls) -> bool: return True @@ -415,7 +426,36 @@ class FlashAttentionImpl(AttentionImpl): "For installation details, see https://gitcode.com/Ascend/MindIE-SD" "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" ) + # Opt-in mask-free paths (mirror the CUDA cu_seqlens behavior): the + # model marks extra["npu_attn_varlen"] and carries packed metadata, so + # the padding document is excluded without reading or materializing + # the quadratic full_qk mask. + # - default: packed varlen via mindiesd attention_forward_varlen; + # - when MINDIE_SD_FA_TYPE=ascend_laser_attention: prefix K/V + # slicing via mindiesd attention_forward (so the env-selected op + # type is honored; the TND varlen op cannot be). + extra = attn_metadata.extra if attn_metadata else {} + if extra.get("npu_attn_varlen", False): + if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention": + out = self._forward_prefix_kv_slice_npu(query, key, value, extra) + else: + out = self._forward_varlen_packed_npu(query, key, value, extra) + if out is not None: + return out attention_mask = attn_metadata.attn_mask if attn_metadata else None + if attention_mask is None and extra.get("npu_attn_varlen", False): + # Models skip mask construction when this opt-in is set + # (FlashAttentionBackend.supports_packed_mask_free). The packed + # paths above declined (contract mismatch), so rebuild the padding + # mask here to keep the masked fallback correct. + used = extra.get("valid_kv_length") + if not isinstance(used, int) or not 0 < used <= query.shape[1]: + raise ValueError( + "npu_attn_varlen packed metadata is unusable and no attn_mask " + f"was constructed (valid_kv_length={used!r}, seq_len={query.shape[1]}); " + "refusing to run unmasked attention over padding rows." + ) + attention_mask = torch.arange(query.shape[1], device=query.device)[None] < used # NPU aclnnFlashAttentionScore requires mask shape to be one of: # [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv] @@ -433,3 +473,161 @@ class FlashAttentionImpl(AttentionImpl): op_type="fused_attn_score", layout=layout, ) + + def _resolve_packed_seq_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + extra: dict, + ) -> tuple[list[int], list[int]] | None: + """Resolve packed document boundaries (cumulative end offsets per doc). + + Returns (seq_q, seq_k), e.g. ([used_q, total_q], [used_k, total_k]), + or None unless the metadata matches the exact contract MiniMax-H3 + produces: + - q/k/v are [1, T, N, D] (BSND, single packed batch); + - cu_seqlens describe a [real, pad] two-document packing, with the + padding document as a strict suffix; + - max_seqlen_q/k are Python ints equal to the real document length + (so boundaries are derived without any device sync). + """ + if self.causal: + return None + packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k") + if any(k not in extra for k in packed_keys): + return None + cu_q, cu_k = extra["cu_seqlens_q"], extra["cu_seqlens_k"] + used_q, used_k = extra["max_seqlen_q"], extra["max_seqlen_k"] + if query.shape[0] != 1 or key.shape[0] != 1: + return None + if not isinstance(used_q, int) or not isinstance(used_k, int): + return None + total_q, total_k = query.shape[1], key.shape[1] + # .shape is host-side metadata: counting documents never syncs. + if cu_q.shape[0] != cu_k.shape[0] or cu_q.shape[0] > 3: + return None + if not (0 < used_q <= total_q) or not (0 < used_k <= total_k): + return None + if used_q == total_q and used_k == total_k: + # No padding document: one full-length document. + return [total_q], [total_k] + if cu_q.shape[0] == 3 and used_q >= total_q - used_q and used_k >= total_k - used_k: + # [real, pad] packing: the real document must be the longer one + # (consistent with the max_seqlen naming). + return [used_q, total_q], [used_k, total_k] + return None + + def _forward_varlen_packed_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + extra: dict, + ) -> torch.Tensor | None: + """Packed varlen attention on NPU via mindiesd attention_forward_varlen. + + Returns None (caller falls back to the mask path) when the packed + contract does not hold; see _resolve_packed_seq_npu. + """ + resolved = self._resolve_packed_seq_npu(query, key, extra) + if resolved is None: + return None + seq_q, seq_k = resolved + + try: + from mindiesd import attention_forward_varlen + except ImportError: + raise ImportError( + "FlashAttentionBackend NPU implementation requires MindIE-SD. " + "Please install MindIE-SD to enable NPU attention support. " + "For installation details, see https://gitcode.com/Ascend/MindIE-SD" + "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" + ) + q = query.squeeze(0) # [T, N, D] == TND + k = key.squeeze(0) + v = value.squeeze(0) + # attention_forward_varlen wants cu_seqlens as host lists of cumulative + # offsets and forwards cu[1:] as actual_seq_qlen/actual_seq_kvlen. + out = attention_forward_varlen( + q, + k, + v, + [0, *seq_q], + [0, *seq_k], + softmax_scale=self.softmax_scale, + ) + return out.unsqueeze(0) + + def _forward_prefix_kv_slice_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + extra: dict, + ) -> torch.Tensor | None: + """Mask-free attention by slicing K/V to the valid prefix (zero-copy + views), then mindiesd attention_forward without a mask. + + Same numerical contract as the varlen path: the padding document is a + strict suffix, so dropping it from K/V is identical to masking it out. + Query keeps full length; outputs on padding rows are never consumed + downstream. Returns None when the packed contract does not hold. + + Used when MINDIE_SD_FA_TYPE=ascend_laser_attention: the TND varlen op + cannot honor the env-selected op type, while attention_forward can. + + The laser kernel stores unscaled S=QK^T in an fp16 GM workspace, so + bf16 activations with large outliers overflow 65504 into ±inf and the + affected rows turn NaN. Models may opt into exact power-of-two input + pre-scaling via extra["laser_input_scale"] (see abstract.py): q/k/v + are divided by the factor, the kernel scale_value is multiplied by its + square, and the output is scaled back. Absent or invalid factor means + no pre-scaling. + """ + resolved = self._resolve_packed_seq_npu(query, key, extra) + if resolved is None: + return None + _, seq_k = resolved + used_k = seq_k[0] # real document length (first cumulative end) + + try: + from mindiesd import attention_forward + except ImportError: + raise ImportError( + "FlashAttentionBackend NPU implementation requires MindIE-SD. " + "Please install MindIE-SD to enable NPU attention support. " + "For installation details, see https://gitcode.com/Ascend/MindIE-SD" + "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" + ) + # mindiesd always takes BSND input regardless of `layout`; the arg + # selects the op-internal layout and laser only supports BNSD, so do + # NOT forward the model's qkv_layout ("BSND" for MiniMax-H3) here. + layout = "BNSD" + key = key[:, :used_k] + value = value[:, :used_k] + + input_scale = extra.get("laser_input_scale") + preserve_input_range = isinstance(input_scale, (int, float)) and input_scale > 1 + if preserve_input_range: + query = query / input_scale + key = key / input_scale + value = value / input_scale + scale = self.softmax_scale + if preserve_input_range: + scale *= input_scale**2 + + def _postprocess(out: torch.Tensor) -> torch.Tensor: + return out * input_scale if preserve_input_range else out + + return _postprocess( + attention_forward( + query, + key, + value, + attn_mask=None, + opt_mode="manual", + op_type="fused_attn_score", # MINDIE_SD_FA_TYPE env overrides this + layout=layout, + scale=scale, + ) + ) diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index d41395bb5..5e180aa23 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -107,6 +107,15 @@ MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"}) # lookup and masked out afterwards). MINIMAX_H3_ADALN_MODALITY_NUM = 3 +# Opt-in fp16-range protection for the NPU ascend_laser_attention kernel +# (consumed only via the "laser_input_scale" extra key; other backends and +# platforms ignore it). The kernel stores unscaled QK^T in an fp16 GM +# workspace, and H3's outlier activations (per-element amax in the hundreds) +# push dot products past fp16 max 65504, turning whole 128-row blocks NaN. +# 256 is a power of two, so pre-dividing q/k/v and the compensating +# kernel-scale/output multiplies are exact in floating point. +MINIMAX_H3_LASER_INPUT_SCALE = 256.0 + def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any: if key not in kwargs or kwargs[key] is None: @@ -402,11 +411,16 @@ class MiniMaxH3Attention(nn.Module): used = min(max_seqlen, packed_total) attn_mask = None # Ring attention can dispatch to a different implementation from the - # configured backend, so this no-mask fast path is local-only. - prefix_slice = ( - not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing + # configured backend, so the no-mask fast paths are local-only. + # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN). + # supports_packed_mask_free: backend consumes the packed metadata + # without ever reading attn_mask (CUDA packed varlen, NPU + # npu_attn_varlen opt-in with its own fallback rebuild). + no_mask = not getattr(self.attention, "use_ring", False) and ( + self.attention.attn_backend.supports_prefix_kv_slicing + or self.attention.attn_backend.supports_packed_mask_free() ) - if used < packed_total and not prefix_slice: + if used < packed_total and not no_mask: attn_mask = torch.arange(packed_total, device=q.device)[None] < used metadata = AttentionMetadata( attn_mask=attn_mask, @@ -416,6 +430,15 @@ class MiniMaxH3Attention(nn.Module): "max_seqlen_q": max_seqlen, "max_seqlen_k": max_seqlen, "valid_kv_length": used, + # Opt the NPU flash backend into the packed varlen path so the + # quadratic full_qk mask is never materialized. Ring attention + # is excluded: it keeps the aligned padding rows for its + # fixed-size P2P buffers and still needs the mask. + "npu_attn_varlen": not getattr(self.attention, "use_ring", False), + # fp16-range protection for the ascend_laser_attention kernel + # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other + # backend/path. + "laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE, }, video_layout=video_layout, ) diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 0421814a0..525725fcf 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -139,6 +139,18 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): ) backend_upper = "FLASH_ATTN" + if backend_upper == "FLASH_ATTN" and find_spec("mindiesd"): + # The NPU FLASH_ATTN backend imports mindiesd lazily at first + # forward, but CANN snapshots the custom-op registry at the + # first custom-op regInfo lookup in the process (e.g. a + # vllm-ascend custom op during model load/warmup). Import + # mindiesd here so its env.py prepends the mindiesd vendor + # dirs (aie_ascendc etc.) to ASCEND_CUSTOM_OPP_PATH before + # that snapshot; otherwise aclnnLaserAttention / + # FusedAttentionScore fail with EZ1001 "does not support + # opType" for the rest of the process. + import mindiesd # noqa: F401 + backend = DiffusionAttentionBackendEnum[backend_upper] logger.debug("Using diffusion attention backend '%s'", backend_upper) return backend.get_path()