init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

vllm_ascend/device/mxfp_compat.py

@@ -0,0 +1,91 @@
+import torch
+import torch_npu
+
+from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type
+
+# TODO(linfeng): Temporary compatibility shim for MXFP4/MXFP8 because current torch_npu
+# releases do not expose the required dtype attributes yet. Simplify or remove this
+# file after the torch_npu release in March 2026 includes those dtype symbols.
+FLOAT8_E8M0FNU_DTYPE = getattr(torch_npu, "float8_e8m0fnu", getattr(torch, "float8_e8m0fnu", None))
+FLOAT4_E2M1FN_X2_DTYPE = getattr(torch_npu, "float4_e2m1fn_x2", getattr(torch, "float4_e2m1fn_x2", None))
+HIFLOAT8_DTYPE = getattr(torch_npu, "hifloat8", None)
+
+
+# TODO(zzzzzz198): Currently three formats(float8_e8m0fnu, float4_e2m1fn_x2, hifloat8) have to be
+# specified for some operators like GMM in Ascend950, while float8_e4m3fn does not. Remove these
+# filterations when operators allow to pass data with these three dtypes directly.
+QUANT_DTYPES = tuple(dtype for dtype in (FLOAT4_E2M1FN_X2_DTYPE, HIFLOAT8_DTYPE) if dtype is not None)
+SCALE_DTYPES = tuple(dtype for dtype in (FLOAT8_E8M0FNU_DTYPE,) if dtype is not None)
+
+
+def _get_missing_symbols(symbols: tuple[str, ...]) -> list[str]:
+    return [symbol for symbol in symbols if not hasattr(torch_npu, symbol)]
+
+
+def _is_dynamic_mx_quant_fusion_soc_supported() -> bool:
+    return get_ascend_device_type() == AscendDeviceType.A5
+
+
+def _ensure_symbols_available(feature: str, symbols: tuple[str, ...]) -> None:
+    missing_symbols = _get_missing_symbols(symbols)
+    if not missing_symbols:
+        return
+    missing_symbols_str = ", ".join(missing_symbols)
+    raise RuntimeError(
+        f"{feature} requires a newer torch_npu runtime. Missing symbols: {missing_symbols_str}. "
+        "Please upgrade torch_npu or disable MXFP quantization."
+    )
+
+
+def is_add_rms_norm_dynamic_mx_quant_fusion_available() -> bool:
+    return (
+        _is_dynamic_mx_quant_fusion_soc_supported()
+        and hasattr(torch, "float8_e4m3fn")
+        and not _get_missing_symbols(("npu_dynamic_mx_quant", "npu_add_rms_norm_dynamic_mx_quant"))
+    )
+
+
+def is_rms_norm_dynamic_mx_quant_fusion_available() -> bool:
+    return (
+        _is_dynamic_mx_quant_fusion_soc_supported()
+        and hasattr(torch, "float8_e4m3fn")
+        and not _get_missing_symbols(("npu_dynamic_mx_quant", "npu_rms_norm_dynamic_mx_quant"))
+    )
+
+
+def ensure_mxfp8_scale_dtype_available(feature: str) -> None:
+    _ensure_symbols_available(feature, ("float8_e8m0fnu",))
+
+
+def ensure_mxfp4_dtype_available(feature: str) -> None:
+    _ensure_symbols_available(feature, ("float4_e2m1fn_x2", "float8_e8m0fnu"))
+
+
+def ensure_mxfp8_linear_available(feature: str) -> None:
+    _ensure_symbols_available(feature, ("float8_e8m0fnu", "npu_dynamic_mx_quant", "npu_quant_matmul"))
+
+
+def ensure_mxfp8_moe_available(feature: str) -> None:
+    _ensure_symbols_available(
+        feature,
+        ("float8_e8m0fnu", "npu_dynamic_mx_quant", "npu_grouped_matmul_swiglu_quant_v2"),
+    )
+
+
+def ensure_mxfp4_linear_available(feature: str) -> None:
+    _ensure_symbols_available(
+        feature, ("float4_e2m1fn_x2", "float8_e8m0fnu", "npu_dynamic_mx_quant", "npu_quant_matmul")
+    )
+
+
+def ensure_mxfp4_flatquant_linear_available(feature: str) -> None:
+    _ensure_symbols_available(
+        feature, ("float4_e2m1fn_x2", "float8_e8m0fnu", "npu_kronecker_quant", "npu_quant_matmul")
+    )
+
+
+def ensure_mxfp4_moe_available(feature: str) -> None:
+    _ensure_symbols_available(
+        feature,
+        ("float4_e2m1fn_x2", "float8_e8m0fnu", "npu_dynamic_mx_quant", "npu_grouped_matmul_swiglu_quant_v2"),
+    )

vllm_ascend/device/utils.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+from typing import Any
+
+import torch
+import torch_npu
+
+FIA_TND_LARGE_HEAD_FALLBACK_HEAD_SIZE = 512
+SWA_INT_MAX = 2147483647
+
+
+def npu_large_head_prefill_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_metadata: Any,
+    *,
+    key_cache: torch.Tensor | None,
+    value_cache: torch.Tensor | None,
+    num_heads: int,
+    num_kv_heads: int,
+    head_size: int,
+    scale: float,
+    is_prefill_no_cache: bool,
+):
+    # A2/A3 FIA TND does not support some large head sizes. Keep those prefill
+    # cases on an NPU attention op instead of falling back to Python.
+    num_tokens = attn_metadata.actual_seq_lengths_q[-1]
+    query = query[:num_tokens]
+    key, value, actual_seq_lengths_kv = _get_large_head_prefill_kv(
+        key,
+        value,
+        attn_metadata,
+        num_tokens,
+        key_cache,
+        value_cache,
+        num_kv_heads,
+        head_size,
+        is_prefill_no_cache,
+    )
+    sparse_mode = 3 if attn_metadata.causal else 0
+    pre_tokens = SWA_INT_MAX
+    next_tokens = 0 if attn_metadata.causal else SWA_INT_MAX
+    attn_mask = attn_metadata.attn_mask
+    if attn_mask is not None and attn_mask.dtype not in (torch.bool, torch.uint8):
+        attn_mask = attn_mask.bool()
+    attn_output = torch_npu.npu_fusion_attention(
+        query=query,
+        key=key,
+        value=value,
+        head_num=num_heads,
+        input_layout="TND",
+        atten_mask=attn_mask,
+        scale=scale,
+        pre_tockens=pre_tokens,
+        next_tockens=next_tokens,
+        actual_seq_qlen=attn_metadata.actual_seq_lengths_q,
+        actual_seq_kvlen=actual_seq_lengths_kv,
+        sparse_mode=sparse_mode,
+    )[0]
+    return attn_output, None
+
+
+def _get_large_head_prefill_kv(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_metadata: Any,
+    num_tokens: int,
+    key_cache: torch.Tensor | None,
+    value_cache: torch.Tensor | None,
+    num_kv_heads: int,
+    head_size: int,
+    is_prefill_no_cache: bool,
+) -> tuple[torch.Tensor, torch.Tensor, list[int]]:
+    # PrefillNoCache already has dense TND key/value tensors. Chunked prefill
+    # may need historical paged KV cache gathered back to dense TND.
+    if is_prefill_no_cache or key_cache is None or value_cache is None:
+        return key[:num_tokens], value[:num_tokens], attn_metadata.actual_seq_lengths_q
+
+    seq_lens = attn_metadata.seq_lens_list
+    if not seq_lens:
+        return key[:num_tokens], value[:num_tokens], attn_metadata.actual_seq_lengths_q
+
+    key, value = _gather_paged_kv_to_dense(
+        key_cache,
+        value_cache,
+        attn_metadata.block_tables,
+        seq_lens,
+        num_kv_heads,
+        head_size,
+    )
+    actual_seq_lengths_kv = []
+    cumsum = 0
+    for length in seq_lens:
+        cumsum += length
+        actual_seq_lengths_kv.append(cumsum)
+    return key, value, actual_seq_lengths_kv
+
+
+def _gather_paged_kv_to_dense(
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    block_table: torch.Tensor,
+    seq_lens: list[int],
+    num_kv_heads: int,
+    head_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # npu_fusion_attention consumes dense TND KV, while cached prefill KV is
+    # stored by blocks. Gather only valid tokens from the block table.
+    block_size = key_cache.shape[1]
+    max_seq_len = max(seq_lens)
+    seq_lens_tensor = torch.tensor(seq_lens, dtype=torch.long, device=key_cache.device)
+    num_blocks = (max_seq_len + block_size - 1) // block_size
+    block_table = block_table[: len(seq_lens), :num_blocks].long()
+
+    flat_block_ids = block_table.reshape(-1)
+    max_tokens_padded = num_blocks * block_size
+    dense_shape = (len(seq_lens), max_tokens_padded, num_kv_heads, head_size)
+    gathered_key = key_cache.index_select(0, flat_block_ids).reshape(dense_shape)
+    gathered_value = value_cache.index_select(0, flat_block_ids).reshape(dense_shape)
+
+    positions = torch.arange(max_tokens_padded, dtype=torch.long, device=key_cache.device)
+    valid_mask = positions.unsqueeze(0) < seq_lens_tensor.unsqueeze(1)
+    return gathered_key[valid_mask].contiguous(), gathered_value[valid_mask].contiguous()