add env to disable qsa kernel

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 "

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

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',
 ]

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,22 +26,42 @@
 """
 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
 except Exception:  # pragma: no cover - triton absent
     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: