diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index 86649e6..03b4768 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -20,8 +20,9 @@ from mcore_bridge.utils import get_env_args, get_local_layer_specs, get_logger from mcore_bridge.utils.megatron_utils import reconstruct_tensor_cp -from ..modules import (GatedDeltaNet, QSAIndexer, QSASparseCoreAttention, Qwen4ExpTextGatedResidual, - Qwen4ExpTextPLELayer, TransformerBlock, TransformerLayer, qsa_sparse_supported) +from ..modules import (QSA_SPARSE_KERNEL_ENV, GatedDeltaNet, QSAIndexer, QSASparseCoreAttention, + Qwen4ExpTextGatedResidual, Qwen4ExpTextPLELayer, TransformerBlock, TransformerLayer, + qsa_sparse_supported, use_qsa_sparse_kernel) from ..register import ModelLoader from .qwen3_next import Qwen3NextBridge, Qwen3NextRMSNorm, Qwen3NextSelfAttention @@ -158,14 +159,20 @@ def _qsa_select(self, hidden_states, attn_kwargs, position_ids=None): if not needs_kernel: return self._qsa_select_mask(hidden_states, attn_kwargs), False - # From here the mask path is not an option, so every failure raises instead of - # silently degrading to dense attention (which would diverge from the sparse - # rollout without telling anyone). + # From here the mask path is not an option, so by default a failure raises instead + # of silently degrading to dense attention (which would diverge from the sparse + # rollout without telling anyone). An explicit opt-out via the env var is the one + # sanctioned fallback: full attention, warned about once. if not sparse_ok: + if not use_qsa_sparse_kernel(): + logger.warning_once(f'QSA sparse kernel is disabled via {QSA_SPARSE_KERNEL_ENV}=0; ' + f'falling back to full attention ({"packing/thd" if is_thd else f"CP={cp_size}"}).') + return None, False raise RuntimeError(f'QSA needs the sparse kernel here ({"packing/thd" if is_thd else f"CP={cp_size}"}), ' 'but QSASparseCoreAttention was not installed -- triton is missing or ' f'kv_channels={getattr(self.config, "kv_channels", None)} is not a power of two. ' - 'Use --padding_free false with context_parallel_size 1 to take the bool-mask path.') + 'Use --padding_free false with context_parallel_size 1 to take the bool-mask path, ' + f'or set {QSA_SPARSE_KERNEL_ENV}=0 to fall back to full attention.') if cp_size > 1 and getattr(self.config, 'cp_comm_type', None) != 'all_gather': raise RuntimeError(f"QSA sparse selection with context_parallel_size={cp_size} requires " f"cp_comm_type='all_gather' (got {getattr(self.config, 'cp_comm_type', None)!r}): the " diff --git a/src/mcore_bridge/model/modules/__init__.py b/src/mcore_bridge/model/modules/__init__.py index fdba46b..996a3bb 100644 --- a/src/mcore_bridge/model/modules/__init__.py +++ b/src/mcore_bridge/model/modules/__init__.py @@ -5,7 +5,7 @@ from .gated_delta_net import GatedDeltaNet from .gated_self_attention import GatedSelfAttention from .hyper_connection_gated import Qwen4ExpTextGatedResidual, Qwen4ExpTextGroupedRMSNorm -from .kernels import QSASparseCoreAttention, qsa_sparse_supported +from .kernels import QSA_SPARSE_KERNEL_ENV, QSASparseCoreAttention, qsa_sparse_supported, use_qsa_sparse_kernel from .mtp_layer import MultiTokenPredictionLayer from .multi_latent_attention import MLASelfAttention from .ple import Qwen4ExpTextNGramEmbedding, Qwen4ExpTextPLELayer diff --git a/src/mcore_bridge/model/modules/kernels/__init__.py b/src/mcore_bridge/model/modules/kernels/__init__.py index 2139173..87f45a2 100644 --- a/src/mcore_bridge/model/modules/kernels/__init__.py +++ b/src/mcore_bridge/model/modules/kernels/__init__.py @@ -1,10 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. from .ple_kernels import gather_ple_rows, ple_gate_conv_triton -from .qsa_kernels import QSASparseCoreAttention, qsa_sparse_supported +from .qsa_kernels import QSA_SPARSE_KERNEL_ENV, QSASparseCoreAttention, qsa_sparse_supported, use_qsa_sparse_kernel __all__ = [ + 'QSA_SPARSE_KERNEL_ENV', 'QSASparseCoreAttention', 'gather_ple_rows', 'ple_gate_conv_triton', 'qsa_sparse_supported', + 'use_qsa_sparse_kernel', ] diff --git a/src/mcore_bridge/model/modules/kernels/qsa_kernels.py b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py index 74821d8..485137c 100644 --- a/src/mcore_bridge/model/modules/kernels/qsa_kernels.py +++ b/src/mcore_bridge/model/modules/kernels/qsa_kernels.py @@ -13,7 +13,8 @@ Context parallelism is blocked for a different reason -- block pooling needs keys from other CP ranks, which the mask path never gathers. - So the layer picks by data shape, with no user-facing switch: + So the layer picks by data shape, unless disabled via the ``QSA_SPARSE_KERNEL`` + env var (then: full attention under thd/CP>1, bool mask otherwise): thd (padding_free) or CP>1 -> this kernel sbhd and CP==1 -> the indexer's bool mask on TE @@ -25,8 +26,14 @@ """ import torch +from mcore_bridge.utils import get_env_args, get_logger + from .qsa_block_sparse_attn import qsa_sparse_attention_from_indices +logger = get_logger() + +QSA_SPARSE_KERNEL_ENV = 'QSA_SPARSE_KERNEL' + try: import triton # noqa: F401 (import guard for the vendored kernel) HAVE_TRITON = True @@ -34,13 +41,27 @@ HAVE_TRITON = False +def use_qsa_sparse_kernel() -> bool: + return get_env_args(QSA_SPARSE_KERNEL_ENV, bool, True) + + def qsa_sparse_supported(head_dim: int) -> bool: """Whether the sparse kernel can run for this head dim. - Triton must be importable and ``head_dim`` must be a power of two (the - kernel tiles the head with ``tl.arange`` blocks). + Triton must be importable, ``head_dim`` must be a power of two (the kernel tiles + the head with ``tl.arange`` blocks), and the kernel must not be disabled via + ``QSA_SPARSE_KERNEL=0``. """ - return HAVE_TRITON and head_dim > 0 and not (head_dim & (head_dim - 1)) + if not use_qsa_sparse_kernel(): + return False + if not HAVE_TRITON or head_dim <= 0 or (head_dim & (head_dim - 1)): + return False + if not torch.cuda.is_available(): + logger.warning_once('The QSA sparse kernel is only tested on CUDA GPUs and may fail to compile on ' + 'this device. If you hit triton compile errors, set ' + f'{QSA_SPARSE_KERNEL_ENV}=0 to disable it (QSA then falls back to full ' + 'attention under packing/CP, and to the bool-mask path otherwise).') + return True def _cp_query_global_positions(seq_len: int, cp_size: int, cp_rank: int, device) -> torch.Tensor: