diff --git a/ATTRIBUTIONS-Python.md b/ATTRIBUTIONS-Python.md index 7e74b846a85f..4d2111f72cdd 100644 --- a/ATTRIBUTIONS-Python.md +++ b/ATTRIBUTIONS-Python.md @@ -5261,7 +5261,7 @@ For more information, please refer to - `Tracker`: https://github.com/tox-dev/py-filelock/issues -## flashinfer-python (0.6.15) +## flashinfer-python (0.6.16) ### Licenses License: `Apache-2.0` diff --git a/requirements.txt b/requirements.txt index e592848fc725..d63bc1a93be2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -58,7 +58,7 @@ ordered-set peft>=0.18.1,<0.19.0 patchelf einops -flashinfer-python==0.6.15 +flashinfer-python==0.6.16 xgrammar==0.1.32 llguidance==0.7.29 jsonschema diff --git a/security_scanning/poetry.lock b/security_scanning/poetry.lock index 084385dc1514..6f9065c5e7fb 100644 --- a/security_scanning/poetry.lock +++ b/security_scanning/poetry.lock @@ -1352,14 +1352,14 @@ dev = ["pytest", "pytest-xdist", "ruff"] [[package]] name = "flashinfer-python" -version = "0.6.15" +version = "0.6.16" description = "FlashInfer: Kernel Library for LLM Serving" optional = false python-versions = "<4.0,>=3.10" groups = ["main"] files = [ - {file = "flashinfer_python-0.6.15-py3-none-any.whl", hash = "sha256:da6c339e14db4831ade0d593324f02907d44bd4b86b640faad4727d9e089477b"}, - {file = "flashinfer_python-0.6.15.tar.gz", hash = "sha256:2a3f1ed47129f9ac9505a26a8f12cadefc0f27d3104fb623ae281032f49eae5f"}, + {file = "flashinfer_python-0.6.16-py3-none-any.whl", hash = "sha256:059d629cf352dae2d2e17c9b632f1be7c12fd7feb4979d71cdffd96104219e53"}, + {file = "flashinfer_python-0.6.16.tar.gz", hash = "sha256:0c89970ee42b0640d3dfd65ea4770710e21101ce2b43a09eda3e29cd1874fe48"}, ] [package.dependencies] @@ -1371,7 +1371,7 @@ einops = "*" nccl4py = ">=0.3.1" ninja = "*" numpy = "*" -nvidia-cudnn-frontend = ">=1.13.0" +nvidia-cudnn-frontend = ">=1.25.0" nvidia-cutlass-dsl = ">=4.5.0" nvidia-ml-py = "*" packaging = ">=24.2" @@ -7187,4 +7187,4 @@ type = ["pytest-mypy (>=1.0.1) ; platform_python_implementation != \"PyPy\""] [metadata] lock-version = "2.1" python-versions = ">=3.10,<3.13" -content-hash = "6efb69c9fb5cf45baaebd398f9dc04301991dad63096035bf05152f009334b3d" +content-hash = "4a5add9776e2a34402b10a4635c87123fc38efb1c580ec4559ae7fb09837a13c" diff --git a/security_scanning/pyproject.toml b/security_scanning/pyproject.toml index 57e011711639..f0f4080fd5a1 100644 --- a/security_scanning/pyproject.toml +++ b/security_scanning/pyproject.toml @@ -55,7 +55,7 @@ dependencies = [ "peft (>=0.18.1,<0.19.0)", "patchelf (>=0.17.2.4,<0.18.0.0)", "einops (>=0.8.2,<0.9.0)", - "flashinfer-python (==0.6.15)", + "flashinfer-python (==0.6.16)", "xgrammar (==0.1.32)", "llguidance (==0.7.29)", "jsonschema (>=4.26.0,<5.0.0)", diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index 26c8045dd89b..51d95634aa7d 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -299,6 +299,7 @@ def _trtllm_gen_batch_decode_with_kv_cache( None, # lse 0, # lse_stride_tokens 0, # lse_stride_heads + False, # enable_block_sparse_attention (added in flashinfer 0.6.16, flashinfer-ai/flashinfer#3955) ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py index 879436f58b2b..6963a046dd68 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py @@ -21,19 +21,49 @@ NVLINK Two-Sided supports post-quant dispatch for all quantization modes. """ +import functools import os from typing import List, Optional, Tuple import pynvml import torch -from flashinfer.comm.mnnvl import MnnvlMemory as flashinfer_MnnvlMemory -from flashinfer.comm.trtllm_alltoall import MnnvlMoe as flashinfer_MnnvlMoe +from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from .base import Communication +@functools.lru_cache(maxsize=1) +def _flashinfer_mnnvl(): + """Resolve FlashInfer's MNNVL symbols on first use. + + Deliberately not a module-level import: ``flashinfer.comm`` does not import + on every interpreter this package supports. FlashInfer >= 0.6.16 annotates + ``array.array[int]`` at module scope in ``comm/fd_exchange.py``, and + ``array.array`` only became subscriptable in Python 3.12, so that annotation + raises ``TypeError`` on the 3.10/3.11 interpreters covered by our + ``python_requires`` (FlashInfer declares ``requires_python >= 3.10`` too, so + this is an upstream bug). At module scope that would break ``import + tensorrt_llm`` outright; resolving it lazily confines the damage to this one + alltoall strategy. + + Returns ``(MnnvlMemory, MnnvlMoe)``, or ``None`` when unavailable. + """ + try: + from flashinfer.comm.mnnvl import MnnvlMemory + from flashinfer.comm.trtllm_alltoall import MnnvlMoe + except (ImportError, TypeError) as e: + # TypeError covers the annotation-evaluation failure described above; + # it is not a generic catch-all. + logger.warning( + f"flashinfer.comm is unavailable ({e}); the NVLinkTwoSidedFlashinfer " + f"MoE alltoall strategy is disabled." + ) + return None + return MnnvlMemory, MnnvlMoe + + class NVLinkTwoSidedFlashinfer(Communication): """ NVLINK two-sided comm AllToAll strategy. @@ -77,10 +107,19 @@ def __init__( # CutlassFusedMoE kernels support any invalid value. self.invalid_token_expert_id: int = -1 + mnnvl = _flashinfer_mnnvl() + if mnnvl is None: + raise RuntimeError( + "NVLinkTwoSidedFlashinfer requires flashinfer.comm, which failed " + "to import (see the preceding warning); select a different MoE " + "alltoall strategy or run on Python >= 3.12." + ) + flashinfer_MnnvlMemory, self._mnnvl_moe = mnnvl + # Initialize NVLINK workspaces flashinfer_MnnvlMemory.initialize() - self.alltoall_workspace = flashinfer_MnnvlMoe.get_moe_workspaces(mapping) - self.alltoall_prepare_workspace = flashinfer_MnnvlMoe.get_moe_prepare_workspace(mapping) + self.alltoall_workspace = self._mnnvl_moe.get_moe_workspaces(mapping) + self.alltoall_prepare_workspace = self._mnnvl_moe.get_moe_prepare_workspace(mapping) # Initialize dispatch state self._dispatch_state = {} @@ -90,6 +129,11 @@ def is_platform_supported() -> bool: """ Check if NVLINK two-sided comm is supported on current hardware. """ + mnnvl = _flashinfer_mnnvl() + if mnnvl is None: + return False + flashinfer_MnnvlMemory, _ = mnnvl + # flashinfer's MnnvlMemory.supports_mnnvl() queries NVML without # initializing it (only its initialize() guards), so satisfy that # precondition on our side of the boundary. @@ -128,7 +172,7 @@ def prepare_dispatch( # Call NVLINK prepare to get alltoall_info and gather EPLB statistics alltoall_info, _, __, gathered_local_statistic_tensor = ( - flashinfer_MnnvlMoe.mnnvl_moe_alltoallv_prepare_without_allgather( + self._mnnvl_moe.mnnvl_moe_alltoallv_prepare_without_allgather( token_selected_slots, None, local_statistic_tensor, @@ -165,7 +209,7 @@ def mnnvl_moe_alltoallv_packed(x, alltoall_info, workspace, ep_rank, ep_size): results = [] for tensor in x: if tensor is not None: - result = flashinfer_MnnvlMoe.mnnvl_moe_alltoallv( + result = self._mnnvl_moe.mnnvl_moe_alltoallv( tensor, alltoall_info, workspace, ep_rank, ep_size ) results.append(result) @@ -221,7 +265,7 @@ def combine( if isinstance(final_hidden_states, list): final_hidden_states = final_hidden_states[0] - final_hidden_states = flashinfer_MnnvlMoe.mnnvl_moe_alltoallv_combine( + final_hidden_states = self._mnnvl_moe.mnnvl_moe_alltoallv_combine( final_hidden_states, self._dispatch_state["alltoall_info"], self.alltoall_workspace,