Compare commits

..
Author SHA1 Message Date
Hongsheng Liu 8e77dae513 docs: fix attention backend link
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-10 21:20:30 +08:00
ooooooye 740f45abcd [Perf][Diffusion] MiniMax-H3: opt into packed varlen attention on NPU to eliminate quadratic mask materialization (#5891)
Signed-off-by: brandneway <gyuan4892@gmail.com>
2026-08-10 19:47:03 +08:00
5e4b9f8c6e [Misc] Align CODEOWNERS with module/feature design docs (#5958)
Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
2026-08-10 19:25:03 +08:00
Hongsheng Liu ae811043e9 docs: keep shared task examples in navigation (#5987)
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-08-10 19:13:52 +08:00
21 changed files with 779 additions and 41 deletions

No files matched your search

+110 -28
View File
@@ -1,6 +1,11 @@
# CODEOWNERS is used to automatically assign PR reviewers # CODEOWNERS is used to automatically assign PR reviewers
# Format: path pattern @username or @team # Format: path pattern @username or @team
# More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners # More info: https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
#
# The vllm_omni/** source-tree rules below are derived from the `owners` and
# `primary_code_paths` front matter in docs/design/module/*.md. When updating
# ownership, update the design doc first, then regenerate the matching rule here.
# Last-match-wins: more specific rules must come AFTER broader ones.
# Default owners (fallback reviewers for all files) # Default owners (fallback reviewers for all files)
* @Gaohan123 @hsliuustc0106 @ywang96 * @Gaohan123 @hsliuustc0106 @ywang96
@@ -21,25 +26,10 @@
# CI # CI
/.buildkite/ @yenuo26 @congw729 @NickCao /.buildkite/ @yenuo26 @congw729 @NickCao
# Engine and model_executor # ComfyUI / external apps
/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 /apps/ @fhfuih @alex-jw-brooks
/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307
/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks @NickCao
/vllm_omni/model_executor/ @tzhouam @gcanlin
/vllm_omni/worker/ @tzhouam @gcanlin
# Diffusion # Benchmarks (out-of-tree)
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
/vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao
# Diffusion CPU offloading
/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi
# Omni and TTS support
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin @NickCao
# Benchmarks
/benchmarks/accuracy/ @david6666666 /benchmarks/accuracy/ @david6666666
/benchmarks/build_dataset/ @linyueqian @yuanheng-zhao /benchmarks/build_dataset/ @linyueqian @yuanheng-zhao
/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi /benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi
@@ -47,20 +37,112 @@
/benchmarks/glm_image/ @princepride @JaredforReal /benchmarks/glm_image/ @princepride @JaredforReal
/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin /benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin
# Usability and Observability # =====================================================================
/vllm_omni/deploy/ @tzhouam @lishunyang12 # Module design ownership (docs/design/module/*.md)
/vllm_omni/metrics/ @lishunyang12 # Each section header cites the source design doc.
/vllm_omni/profiler/ @gcanlin # =====================================================================
# ComfyUI # Autoregressive runtime — docs/design/module/ar_runtime.md
/apps/ @fhfuih @alex-jw-brooks /vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
/vllm_omni/worker/ @tzhouam @yinpeiqi @fake0fan @Sy0307 @Gaohan123
# Large Scale Serving # Engine orchestration & stage runtime — docs/design/module/engine_orchestration.md, stage_runtime.md
/vllm_omni/engine/ @tzhouam @fake0fan
# File-level overrides: input/output modality contracts own these engine files
# — docs/design/module/input_output_modality_contracts.md
/vllm_omni/engine/messages.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/serialization.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/mm_outputs.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/output_modality.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/engine/output_processor.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
# Input/output modality contracts — docs/design/module/input_output_modality_contracts.md
/vllm_omni/inputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/outputs/ @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/request.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
/vllm_omni/data_entry_keys.py @Sy0307 @amy-why-3459 @Gaohan123 @alex-jw-brooks
# Entrypoints and serving boundaries — docs/design/module/entrypoints.md
/vllm_omni/entrypoints/ @alex-jw-brooks @linyueqian @NickCao
# File-level overrides within entrypoints/
/vllm_omni/entrypoints/openai/errors.py @alex-jw-brooks @NickCao
/vllm_omni/entrypoints/cli/benchmark/ @alex-jw-brooks @Bounty-hunter
# Error contracts — docs/design/module/error_contracts.md
/vllm_omni/errors.py @alex-jw-brooks @NickCao
# OmniConnector (inter-stage transport) — docs/design/module/omni_connector.md
/vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan /vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan
# Quantization and configuration # Model integration — docs/design/module/model_integration.md
/vllm_omni/model_executor/ @tzhouam @gcanlin
/vllm_omni/model_executor/models/ @tzhouam @gcanlin @Sy0307 @amy-why-3459 @linyueqian
/vllm_omni/model_extras/ @tzhouam @gcanlin @princepride
/vllm_omni/plugins/ @gcanlin @tjtanaa @xuechendi
# Config override within model_executor — docs/design/module/vllm_omni_config.md
/vllm_omni/model_executor/stage_configs/ @lishunyang12 @alex-jw-brooks
# Execution platforms (hardware backends) — docs/design/module/execution_platforms.md
/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi
/vllm_omni/attention/ @gcanlin @tjtanaa @xuechendi @david6666666 @lishunyang12
# Configuration & deploy — docs/design/module/vllm_omni_config.md
/vllm_omni/config/ @lishunyang12 @alex-jw-brooks /vllm_omni/config/ @lishunyang12 @alex-jw-brooks
/vllm_omni/deploy/ @lishunyang12 @alex-jw-brooks
# Quantization — docs/design/module/quantization.md
/vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12 /vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12
# Hardware plugins and accelerators # Observability — docs/design/module/observability.md
/vllm_omni/platforms/ @gcanlin @tjtanaa @xuechendi /vllm_omni/metrics/ @lishunyang12 @vraiti
# Profiling — docs/design/module/profiling.md
/vllm_omni/profiler/ @gcanlin
# Benchmarking (in-tree) — docs/design/module/benchmarking.md
/vllm_omni/benchmarks/ @alex-jw-brooks @Bounty-hunter
# LoRA (no dedicated design doc; ownership intentionally narrowed)
/vllm_omni/lora/ @SamitHuang @wtomin
# Transformers utils (related path only in design docs; preserves prior ownership)
/vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao
# =====================================================================
# Diffusion module family — docs/design/module/diffusion/
# Umbrella rule first; sub-module rules after it win via last-match.
# =====================================================================
# Diffusion umbrella (top-level files: data.py, forward_context.py, etc.)
# — docs/design/module/diffusion/index.md
/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi
# Cache management — docs/design/module/cache_management.md
/vllm_omni/diffusion/cache/ @Isotr0py @princepride @SamitHuang
# Diffusion quantization — docs/design/module/quantization.md
/vllm_omni/diffusion/quantization/ @david6666666 @Isotr0py @lishunyang12
# Diffusion profiling — docs/design/module/profiling.md
/vllm_omni/diffusion/profiler/ @gcanlin
# Diffusion runtime (scheduler, executor, worker)
# — docs/design/module/diffusion/diffusion_runtime.md, continuous_batching.md
/vllm_omni/diffusion/sched/ @Isotr0py @princepride @SamitHuang @fhfuih
/vllm_omni/diffusion/executor/ @Isotr0py @princepride @SamitHuang @fhfuih
/vllm_omni/diffusion/worker/ @Isotr0py @princepride @SamitHuang @fhfuih
# Diffusion model integration (pipelines, registration, checkpoint loading)
# — docs/design/module/diffusion/diffusion_model_integration.md
/vllm_omni/diffusion/models/ @fhfuih @Bounty-hunter @wtomin
/vllm_omni/diffusion/model_loader/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
/vllm_omni/diffusion/lora/ @fhfuih @Bounty-hunter @wtomin @RuixiangMa
# Diffusion parallelism (rank topology, TP/SP/PP/VAE patch)
# — docs/design/module/diffusion/parallelism.md
/vllm_omni/diffusion/distributed/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
/vllm_omni/diffusion/attention/parallel/ @princepride @yuanheng-zhao @xuechendi @Bounty-hunter
# Diffusion offloader (CPU offloading, DLO)
# — docs/design/module/diffusion/offloader.md
/vllm_omni/diffusion/offloader/ @yuanheng-zhao @david6666666 @lishunyang12
+1
View File
@@ -7,6 +7,7 @@ owners:
- "@yinpeiqi" - "@yinpeiqi"
- "@fake0fan" - "@fake0fan"
- "@Sy0307" - "@Sy0307"
- "@Gaohan123"
primary_code_paths: primary_code_paths:
- vllm_omni/core/** - vllm_omni/core/**
- vllm_omni/worker/** - vllm_omni/worker/**
+1
View File
@@ -4,6 +4,7 @@ kind: module
status: draft status: draft
owners: owners:
- "@alex-jw-brooks" - "@alex-jw-brooks"
- "@Bounty-hunter"
primary_code_paths: primary_code_paths:
- vllm_omni/benchmarks/** - vllm_omni/benchmarks/**
- vllm_omni/entrypoints/cli/benchmark/** - vllm_omni/entrypoints/cli/benchmark/**
+1
View File
@@ -5,6 +5,7 @@ status: draft
owners: owners:
- "@Isotr0py" - "@Isotr0py"
- "@princepride" - "@princepride"
- "@SamitHuang"
primary_code_paths: primary_code_paths:
- vllm_omni/diffusion/cache/** - vllm_omni/diffusion/cache/**
related_code_paths: related_code_paths:
@@ -3,9 +3,10 @@ title: Diffusion Model Integration
kind: module kind: module
status: draft status: draft
owners: owners:
- "@Isotr0py" - "@fhfuih"
- "@Bounty-hunter"
- "@wtomin" - "@wtomin"
- "@ZJY0516" - "@RuixiangMa"
primary_code_paths: primary_code_paths:
- vllm_omni/diffusion/models/** - vllm_omni/diffusion/models/**
- vllm_omni/diffusion/model_loader/** - vllm_omni/diffusion/model_loader/**
@@ -6,6 +6,7 @@ owners:
- "@Isotr0py" - "@Isotr0py"
- "@princepride" - "@princepride"
- "@SamitHuang" - "@SamitHuang"
- "@fhfuih"
primary_code_paths: primary_code_paths:
- vllm_omni/diffusion/sched/** - vllm_omni/diffusion/sched/**
- vllm_omni/diffusion/executor/** - vllm_omni/diffusion/executor/**
+2 -1
View File
@@ -4,7 +4,8 @@ kind: module
status: draft status: draft
owners: owners:
- "@yuanheng-zhao" - "@yuanheng-zhao"
- "@Isotr0py" - "@david6666666"
- "@lishunyang12"
primary_code_paths: primary_code_paths:
- vllm_omni/diffusion/offloader/** - vllm_omni/diffusion/offloader/**
related_code_paths: related_code_paths:
@@ -6,6 +6,7 @@ owners:
- "@princepride" - "@princepride"
- "@yuanheng-zhao" - "@yuanheng-zhao"
- "@xuechendi" - "@xuechendi"
- "@Bounty-hunter"
primary_code_paths: primary_code_paths:
- vllm_omni/diffusion/distributed/** - vllm_omni/diffusion/distributed/**
- vllm_omni/diffusion/attention/parallel/** - vllm_omni/diffusion/attention/parallel/**
+1
View File
@@ -6,6 +6,7 @@ architecture_state: current-plus-in-flight
owners: owners:
- "@alex-jw-brooks" - "@alex-jw-brooks"
- "@linyueqian" - "@linyueqian"
- "@NickCao"
document_stewards: document_stewards:
- "@hsliuustc0106" - "@hsliuustc0106"
- "@Gaohan123" - "@Gaohan123"
+8 -1
View File
@@ -8,8 +8,15 @@ owners:
- "@xuechendi" - "@xuechendi"
primary_code_paths: primary_code_paths:
- vllm_omni/platforms/** - vllm_omni/platforms/**
related_code_paths:
- vllm_omni/attention/** - vllm_omni/attention/**
primary_code_path_owners:
# Additional owners scoped to attention kernels (quantized/custom attention)
- paths:
- vllm_omni/attention/**
owners:
- "@david6666666"
- "@lishunyang12"
related_code_paths: []
depends_on: depends_on:
- ar_runtime.md - ar_runtime.md
- diffusion/index.md - diffusion/index.md
@@ -6,6 +6,8 @@ architecture_state: current
owners: owners:
- "@Sy0307" - "@Sy0307"
- "@amy-why-3459" - "@amy-why-3459"
- "@Gaohan123"
- "@alex-jw-brooks"
document_stewards: document_stewards:
- "@hsliuustc0106" - "@hsliuustc0106"
- "@Gaohan123" - "@Gaohan123"
+8
View File
@@ -9,6 +9,14 @@ primary_code_paths:
- vllm_omni/model_executor/** - vllm_omni/model_executor/**
- vllm_omni/model_extras/** - vllm_omni/model_extras/**
- vllm_omni/plugins/** - vllm_omni/plugins/**
primary_code_path_owners:
# Additional owners scoped to a sub-path of model_executor/**
- paths:
- vllm_omni/model_executor/models/**
owners:
- "@Sy0307"
- "@amy-why-3459"
- "@linyueqian"
related_code_paths: related_code_paths:
- vllm_omni/transformers_utils/** - vllm_omni/transformers_utils/**
- vllm_omni/tokenizers/** - vllm_omni/tokenizers/**
+1
View File
@@ -4,6 +4,7 @@ kind: module
status: draft status: draft
owners: owners:
- "@lishunyang12" - "@lishunyang12"
- "@vraiti"
primary_code_paths: primary_code_paths:
- vllm_omni/metrics/** - vllm_omni/metrics/**
related_code_paths: related_code_paths:
@@ -99,7 +99,7 @@ a native vLLM-Omni pipeline for those cases.
### Attention Backends ### Attention Backends
The diffusers backend converts The diffusers backend converts
[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md) [vLLM-Omni standard of attention backend setting](../../diffusion/attention_backends.md)
to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends). to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).
Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2. Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.
+67 -4
View File
@@ -79,7 +79,8 @@ sharing the tokenizer, processor, text encoder, and VAEs from `FL2VA`.
### Multi-NPU: 768P combined-service configuration ### Multi-NPU: 768P combined-service configuration
Use Ulysses sequence parallelism degree 8, text-encoder tensor parallelism 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 ```bash
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
@@ -98,7 +99,7 @@ vllm serve "${MODEL}" \
--usp 8 \ --usp 8 \
--ring 1 \ --ring 1 \
--text-encoder-tp-size 8 \ --text-encoder-tp-size 8 \
--enable-layerwise-offload \ --enable-distributed-layerwise-offload \
--vae-parallel-mode tile \ --vae-parallel-mode tile \
--vae-use-tiling \ --vae-use-tiling \
--vae-patch-parallel-size 8 \ --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 The same endpoint accepts `task=t2va`, `task=fl2va`, and `task=ref2va`; no
partition restart is required. Layerwise offload applies to both DiTs. 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 ### Optional optimizations
Two independent optimizations may be enabled on top of the configuration 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 over the whole sequence, so ring parallelism would split away the keys it
needs. Scale with `--usp` instead. 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 ## HTTP API examples
The request format is identical to the GPU recipe; see The request format is identical to the GPU recipe; see
@@ -154,8 +217,8 @@ configuration above:
| Workload | Configuration | | Workload | Configuration |
|----------|---------------| |----------|---------------|
| T2VA, 209 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, 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 These measurements describe the validated shapes rather than a general
throughput guarantee. throughput guarantee.
@@ -9,6 +9,10 @@ This script tests two main scenarios:
2. Case 2: Comparing FlashAttention and SDPA backends for batch_size=2 with padding 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 pytest
import torch import torch
@@ -371,6 +375,264 @@ def test_packed_varlen_metadata_must_be_complete(monkeypatch):
impl.forward_cuda(query, query, query, metadata) 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__": if __name__ == "__main__":
print("Running FlashAttention Padding Tests...") print("Running FlashAttention Padding Tests...")
print("=" * 60) print("=" * 60)
@@ -492,6 +492,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
class FakeBackend: class FakeBackend:
supports_prefix_kv_slicing = True supports_prefix_kv_slicing = True
@classmethod
def supports_packed_mask_free(cls) -> bool:
return False
class FakeAttention(torch.nn.Module): class FakeAttention(torch.nn.Module):
attn_backend = FakeBackend 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 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(): def test_packed_attention_keeps_padding_mask_for_other_backends():
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
MiniMaxH3Attention, MiniMaxH3Attention,
@@ -534,6 +580,10 @@ def test_packed_attention_keeps_padding_mask_for_other_backends():
class FakeBackend: class FakeBackend:
supports_prefix_kv_slicing = False supports_prefix_kv_slicing = False
@classmethod
def supports_packed_mask_free(cls) -> bool:
return False
class FakeAttention(torch.nn.Module): class FakeAttention(torch.nn.Module):
attn_backend = FakeBackend attn_backend = FakeBackend
@@ -20,6 +20,17 @@ class AttentionBackend(ABC):
# avoid a slower masked-attention plan when tail padding is not semantic. # avoid a slower masked-attention plan when tail padding is not semantic.
supports_prefix_kv_slicing: bool = False 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. # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted.
# Platform resolution rejects an explicit selection outside this set, so a # Platform resolution rejects an explicit selection outside this set, so a
# hardware-specific backend fails before the model is built. # hardware-specific backend fails before the model is built.
@@ -121,6 +132,17 @@ class AttentionMetadata:
# the packed cu_seqlens tensors. # the packed cu_seqlens tensors.
# "valid_kv_length": int — contiguous valid K/V prefix length for a # "valid_kv_length": int — contiguous valid K/V prefix length for a
# backend that advertises supports_prefix_kv_slicing. # 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). # Piecewise attention metadata (mixed causal/full masks).
# full_attn_spans: per-sample [start, end) spans in global coordinates using full attention. # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention.
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0 # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
from functools import partial from functools import partial
import torch import torch
@@ -16,6 +17,7 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
piecewise_attn, piecewise_attn,
) )
from vllm_omni.diffusion.config import get_current_diffusion_config_or_none from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
from vllm_omni.platforms import current_omni_platform
logger = init_logger(__name__) logger = init_logger(__name__)
@@ -24,6 +26,15 @@ class FlashAttentionBackend(AttentionBackend):
accept_output_buffer: bool = True accept_output_buffer: bool = True
supports_piecewise_spans: 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 @classmethod
def supports_attention_mask(cls) -> bool: def supports_attention_mask(cls) -> bool:
return True return True
@@ -415,7 +426,36 @@ class FlashAttentionImpl(AttentionImpl):
"For installation details, see https://gitcode.com/Ascend/MindIE-SD" "For installation details, see https://gitcode.com/Ascend/MindIE-SD"
"Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" "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 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: # NPU aclnnFlashAttentionScore requires mask shape to be one of:
# [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv] # [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", op_type="fused_attn_score",
layout=layout, 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,
)
)
@@ -107,6 +107,15 @@ MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"})
# lookup and masked out afterwards). # lookup and masked out afterwards).
MINIMAX_H3_ADALN_MODALITY_NUM = 3 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: def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any:
if key not in kwargs or kwargs[key] is None: if key not in kwargs or kwargs[key] is None:
@@ -402,11 +411,16 @@ class MiniMaxH3Attention(nn.Module):
used = min(max_seqlen, packed_total) used = min(max_seqlen, packed_total)
attn_mask = None attn_mask = None
# Ring attention can dispatch to a different implementation from the # Ring attention can dispatch to a different implementation from the
# configured backend, so this no-mask fast path is local-only. # configured backend, so the no-mask fast paths are local-only.
prefix_slice = ( # supports_prefix_kv_slicing: backend slices K/V itself (cuDNN).
not getattr(self.attention, "use_ring", False) and self.attention.attn_backend.supports_prefix_kv_slicing # 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 attn_mask = torch.arange(packed_total, device=q.device)[None] < used
metadata = AttentionMetadata( metadata = AttentionMetadata(
attn_mask=attn_mask, attn_mask=attn_mask,
@@ -416,6 +430,15 @@ class MiniMaxH3Attention(nn.Module):
"max_seqlen_q": max_seqlen, "max_seqlen_q": max_seqlen,
"max_seqlen_k": max_seqlen, "max_seqlen_k": max_seqlen,
"valid_kv_length": used, "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, video_layout=video_layout,
) )
+12
View File
@@ -139,6 +139,18 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
) )
backend_upper = "FLASH_ATTN" 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] backend = DiffusionAttentionBackendEnum[backend_upper]
logger.debug("Using diffusion attention backend '%s'", backend_upper) logger.debug("Using diffusion attention backend '%s'", backend_upper)
return backend.get_path() return backend.get_path()