Batch invariant support PART1

lmdeploy/cli/serve.py

@@ -99,6 +99,7 @@ def add_parser_api_server():
         ArgumentHelper.adapters(pt_group)
         ArgumentHelper.device(pt_group)
         ArgumentHelper.eager_mode(pt_group)
+        ArgumentHelper.enable_batch_invariant(pt_group)
         ArgumentHelper.logprobs_mode(pt_group)
         ArgumentHelper.dllm_block_length(pt_group)
         ArgumentHelper.dllm_unmasking_strategy(pt_group)
@@ -236,6 +237,7 @@ def api_server(args):
                 device_type=args.device,
                 quant_policy=args.quant_policy,
                 eager_mode=args.eager_mode,
+                enable_batch_invariant=args.enable_batch_invariant,
                 max_prefill_token_num=args.max_prefill_token_num,
                 enable_microbatch=args.enable_microbatch,
                 enable_eplb=args.enable_eplb,

lmdeploy/cli/utils.py

@@ -646,6 +646,18 @@ def eager_mode(parser):
                                    help='Whether to enable eager mode. '
                                    'If True, cuda graph would be disabled')
 
+    @staticmethod
+    def enable_batch_invariant(parser):
+        """Add argument enable_batch_invariant to parser."""
+
+        return parser.add_argument('--enable-batch-invariant',
+                                   action='store_true',
+                                   default=False,
+                                   help='Enable batch-invariant greedy inference for supported PyTorch CUDA Hopper '
+                                   'runtime shapes. Warning: expert parallelism/DeepEP is unsupported, and active '
+                                   'eviction/recompute may break batch-invariant behavior. Can also be enabled with '
+                                   'LMDEPLOY_ENABLE_BATCH_INVARIANT=1.')
+
     @staticmethod
     def communicator(parser):
         return parser.add_argument('--communicator',

lmdeploy/messages.py

@@ -366,6 +366,12 @@ class PytorchEngineConfig:
         enable_prefix_caching: Enable token match and sharing caches.
         device_type: The inference device type, options ['cuda']
         eager_mode: Enable "eager" mode or not
+        enable_batch_invariant: Enable batch-invariant greedy inference on
+            supported CUDA backends. The initial implementation is Hopper-only.
+            It can also be enabled process-wide with
+            `LMDEPLOY_ENABLE_BATCH_INVARIANT=1`. Warning: expert
+            parallelism/DeepEP is unsupported, and active eviction/recompute
+            may break batch-invariant behavior.
         custom_module_map: nn module map customized by users. Once
             provided, the original nn modules of the model will be
             substituted by the mapping ones
@@ -426,6 +432,7 @@ class PytorchEngineConfig:
     enable_prefix_caching: bool = False
     device_type: str = 'cuda'
     eager_mode: bool = False
+    enable_batch_invariant: bool = False
     custom_module_map: dict[str, str] = None
     download_dir: str = None
     revision: str = None
@@ -473,6 +480,8 @@ def __post_init__(self):
         except ValueError as e:
             raise ValueError(f'invalid quant_policy: {self.quant_policy}') from e
         assert self.device_type in ['cuda', 'ascend', 'maca', 'camb'], (f'invalid device_type: {self.device_type}')
+        if self.enable_batch_invariant and self.device_type != 'cuda':
+            raise ValueError('enable_batch_invariant is currently supported only for CUDA backend.')
         assert self.kernel_block_size >= 16 and \
                (self.kernel_block_size & (self.kernel_block_size - 1)) == 0, \
                f'kernel_block_size must be >= 16 and a power of 2, but got {self.kernel_block_size}'

lmdeploy/pytorch/backends/base.py

@@ -108,3 +108,9 @@ def device_count():
     def support_ray():
         """Support ray."""
         return False
+
+    @staticmethod
+    def apply_backend_policy(backend_config: BackendConfig, validate_device: bool = False):
+        """Apply backend-specific runtime policy."""
+        if backend_config.enable_batch_invariant:
+            raise RuntimeError('Selected backend does not support enable_batch_invariant.')

lmdeploy/pytorch/backends/cuda/attention/__init__.py

@@ -32,6 +32,22 @@ def use_fa3_warning():
     return False
 
 
+def _fa3_disabled_reasons(alibi: bool, learnable_sink: bool, block_sparse_size: int, head_size: int):
+    """Get reasons why FA3 attention cannot be enabled."""
+    reasons = []
+    if alibi:
+        reasons.append('alibi=True')
+    if learnable_sink:
+        reasons.append('learnable_sink=True')
+    if block_sparse_size != 1:
+        reasons.append(f'block_sparse_size={block_sparse_size}')
+    if head_size > 256:
+        reasons.append(f'head_size={head_size}')
+    if not use_fa3:
+        reasons.append('FlashAttention-3 is unavailable')
+    return reasons
+
+
 @functools.lru_cache
 def _enable_fa3(alibi: bool, learnable_sink: bool, block_sparse_size: int, head_size: int) -> bool:
     """Check if FA3 should be enabled.
@@ -130,6 +146,14 @@ def build(
         )
         enable_fa3 = _enable_fa3(alibi, learnable_sink, block_sparse_size, head_size)
 
+        from ..batch_invariant import is_batch_invariant_policy_enabled
+        if is_batch_invariant_policy_enabled():
+            if use_flash_mla:
+                raise RuntimeError('enable_batch_invariant requires FA3 attention, but this layer uses FlashMLA.')
+            if not enable_fa3:
+                reasons = ', '.join(_fa3_disabled_reasons(alibi, learnable_sink, block_sparse_size, head_size))
+                raise RuntimeError(f'enable_batch_invariant requires FA3 attention, but FA3 is disabled: {reasons}.')
+
         if use_flash_mla is True:
             logger.debug('Build FlashMLAImpl Attention')
             from .mla import FlashMLAImpl

lmdeploy/pytorch/backends/cuda/attention/fa3.py

@@ -19,7 +19,7 @@ class FA3Impl(TritonAttentionImpl):
     Key features:
     - Optimized prefill using flash_attn_varlen_func
     - Speculative decoding support with multi-token queries
-    - Standard single-token decoding with paged attention
+    - Standard single-token decoding with flash_attn_with_kvcache
     """
 
     def __init__(
@@ -79,6 +79,28 @@ def _normalize_sliding_window(self, sliding_window):
             return (sliding_window, sliding_window)
         return sliding_window
 
+    @staticmethod
+    def _decode_num_splits() -> int:
+        """Get FA3 decode split policy."""
+        from lmdeploy.pytorch.backends.cuda.batch_invariant import get_fa3_decode_num_splits
+
+        return get_fa3_decode_num_splits()
+
+    def _raise_if_batch_invariant_quantized_decode(self, quant_policy: QuantPolicy):
+        """Reject quantized KV decode in batch-invariant mode."""
+        if quant_policy == QuantPolicy.NONE or self._decode_num_splits() != 1:
+            return
+        raise RuntimeError(
+            'enable_batch_invariant does not support quantized KV cache with FA3 attention. '
+            'Quantized KV decode would fall back to Triton paged attention or raw FA3 cache access.')
+
+    @staticmethod
+    def _get_fa3_cache_seqlens(attn_metadata: TritonAttentionMetadata) -> torch.Tensor:
+        """Get FA3-compatible cache sequence lengths."""
+        if attn_metadata.kv_seqlens.dtype != torch.int32:
+            attn_metadata.kv_seqlens = attn_metadata.kv_seqlens.to(torch.int32)
+        return attn_metadata.kv_seqlens
+
     def _decoding_speculative(
         self,
         query: torch.Tensor,
@@ -104,6 +126,7 @@ def _decoding_speculative(
             Attention output tensor.
         """
         quant_policy = attn_metadata.quant_policy
+        self._raise_if_batch_invariant_quantized_decode(quant_policy)
 
         # TurboQuant stores packed uint8 data in cache, which FA3's native
         # flash_attn_with_kvcache cannot dequantize directly.
@@ -126,14 +149,15 @@ def _decoding_speculative(
             query,
             k_cache,
             v_cache,
-            cache_seqlens=attn_metadata.kv_seqlens.to(torch.int32),
+            cache_seqlens=self._get_fa3_cache_seqlens(attn_metadata),
             max_seqlen_q=max_q_seqlen,
             scheduler_metadata=attn_metadata.scheduler_metadata,
             page_table=block_offsets,
             softmax_scale=self.scale,
             causal=self.causal,
             window_size=sliding_window,
             softcap=self.logit_softcapping,
+            num_splits=self._decode_num_splits(),
         )
         return attn_output
 
@@ -150,7 +174,8 @@ def _decoding_standard(
         """Standard single-token decoding.
 
         This path handles standard decoding where only one token is generated
-        per request (max_q_seqlen = 1). Uses paged attention for memory efficiency.
+        per request (max_q_seqlen = 1). Uses FA3's flash_attn_with_kvcache unless
+        the KV cache is quantized.
 
         Args:
             query: Query tensor (single token per request).
@@ -166,26 +191,47 @@ def _decoding_standard(
         """
         block_offsets = attn_metadata.block_offsets
         quant_policy = attn_metadata.quant_policy
+        self._raise_if_batch_invariant_quantized_decode(quant_policy)
+        sliding_window = self._normalize_sliding_window(self.sliding_window)
+
+        if quant_policy != QuantPolicy.NONE:
+            attn_output = self.paged_attention_fwd(
+                query,
+                k_cache,
+                v_cache,
+                cache_seqlens=attn_metadata.kv_seqlens,
+                page_table=block_offsets,
+                cu_seqlens_q=attn_metadata.cu_seqlens_q,
+                max_seqlen_q=max_q_seqlen,
+                scheduler_metadata=attn_metadata.scheduler_metadata,
+                softmax_scale=self.scale,
+                causal=self.causal,
+                softcap=self.logit_softcapping,
+                window_size=self.sliding_window,
+                # custom args
+                k_scales_zeros=k_scales_zeros,
+                v_scales_zeros=v_scales_zeros,
+                quant_policy=quant_policy,
+            )
+            return attn_output
 
-        attn_output = self.paged_attention_fwd(
+        query_shape = query.shape
+        query = query.unflatten(0, (-1, max_q_seqlen))
+        attn_output = self.flash_attn_with_kvcache_v3(
             query,
             k_cache,
             v_cache,
-            cache_seqlens=attn_metadata.kv_seqlens,
-            page_table=block_offsets,
-            cu_seqlens_q=attn_metadata.cu_seqlens_q,
+            cache_seqlens=self._get_fa3_cache_seqlens(attn_metadata),
             max_seqlen_q=max_q_seqlen,
             scheduler_metadata=attn_metadata.scheduler_metadata,
+            page_table=block_offsets,
             softmax_scale=self.scale,
             causal=self.causal,
+            window_size=sliding_window,
             softcap=self.logit_softcapping,
-            window_size=self.sliding_window,
-            # custom args
-            k_scales_zeros=k_scales_zeros,
-            v_scales_zeros=v_scales_zeros,
-            quant_policy=quant_policy,
+            num_splits=self._decode_num_splits(),
         )
-        return attn_output
+        return attn_output.reshape(query_shape)
 
     def _forward_decoding(
         self,

lmdeploy/pytorch/backends/cuda/batch_invariant.py

@@ -0,0 +1,124 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+"""Batch-invariant CUDA math policy helpers."""
+
+import os
+
+import torch
+
+from lmdeploy.pytorch.config import BackendConfig
+from lmdeploy.utils import get_logger
+
+logger = get_logger('lmdeploy')
+
+_CUBLAS_WORKSPACE_CONFIG = ':16:8'
+_CUBLASLT_WORKSPACE_SIZE = '1'
+_BATCH_INVARIANT_POLICY_ENABLED = False
+
+
+def is_batch_invariant_policy_enabled() -> bool:
+    """Return whether the CUDA batch-invariant policy has been applied."""
+    return _BATCH_INVARIANT_POLICY_ENABLED
+
+
+def get_fa3_decode_num_splits() -> int:
+    """Get FA3 decode split policy."""
+    return 1 if is_batch_invariant_policy_enabled() else 0
+
+
+def _cuda_initialized() -> bool:
+    """Return whether CUDA is already initialized without initializing it."""
+    try:
+        return torch.cuda.is_initialized()
+    except Exception:
+        return False
+
+
+def _set_env_before_cuda_init(name: str, value: str):
+    """Set an env var that must take effect before CUDA/cuBLAS init."""
+    current = os.environ.get(name)
+    if current == value:
+        return
+    if current is not None and _cuda_initialized():
+        raise RuntimeError(
+            f'enable_batch_invariant requires {name}={value}, but found {name}={current} '
+            'after CUDA was already initialized. Set the environment before importing/running LMDeploy, '
+            'or create the engine before any CUDA work in this process.')
+    if current is not None:
+        logger.warning(f'Override {name}={current} with {name}={value} for enable_batch_invariant.')
+    os.environ[name] = value
+
+
+def _set_torch_precision_policy():
+    """Disable CUDA matmul precision shortcuts that can change reductions."""
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    torch.set_float32_matmul_precision('highest')
+
+    if hasattr(torch.backends.cuda.matmul, 'fp32_precision'):
+        torch.backends.cuda.matmul.fp32_precision = 'ieee'
+    cudnn_conv = getattr(torch.backends.cudnn, 'conv', None)
+    if cudnn_conv is not None and hasattr(cudnn_conv, 'fp32_precision'):
+        cudnn_conv.fp32_precision = 'ieee'
+    cudnn_rnn = getattr(torch.backends.cudnn, 'rnn', None)
+    if cudnn_rnn is not None and hasattr(cudnn_rnn, 'fp32_precision'):
+        cudnn_rnn.fp32_precision = 'ieee'
+
+    preferred_blas_library = getattr(torch.backends.cuda, 'preferred_blas_library', None)
+    if preferred_blas_library is not None:
+        preferred_blas_library(backend='cublaslt')
+
+    matmul_backend = torch.backends.cuda.matmul
+    for dtype_name in ('fp16', 'bf16'):
+        attr_name = f'allow_{dtype_name}_reduced_precision_reduction'
+        if hasattr(matmul_backend, attr_name):
+            try:
+                setattr(matmul_backend, attr_name, (False, False))
+            except TypeError:
+                setattr(matmul_backend, attr_name, False)
+        split_k_attr_name = f'{attr_name}_split_k'
+        if hasattr(matmul_backend, split_k_attr_name):
+            try:
+                setattr(matmul_backend, split_k_attr_name, False)
+            except AttributeError:
+                pass
+
+
+def validate_batch_invariant_device(device: int | torch.device | None = None):
+    """Validate that the current CUDA device is in the initial support
+    scope."""
+    if not torch.cuda.is_available():
+        raise RuntimeError('enable_batch_invariant currently requires a CUDA Hopper GPU.')
+
+    major, minor = torch.cuda.get_device_capability(device)
+    if major != 9:
+        device_name = torch.cuda.get_device_name(device)
+        raise RuntimeError(
+            'enable_batch_invariant currently supports Hopper CUDA GPUs only, '
+            f'but detected {device_name} with SM{major}.{minor}.')
+
+
+def apply_batch_invariant_policy(config: BackendConfig,
+                                 *,
+                                 validate_device: bool = False):
+    """Apply the opt-in Hopper batch-invariant policy.
+
+    The environment settings must be applied before CUDA/cuBLAS/cuBLASLt choose algorithms. Call this once in the parent
+    process before any CUDA query, and again in worker processes before they build models or initialize CUDA-heavy
+    runtime state.
+    """
+    global _BATCH_INVARIANT_POLICY_ENABLED
+
+    if not config.enable_batch_invariant:
+        return
+
+    if config.device_type != 'cuda':
+        raise ValueError('enable_batch_invariant is currently supported only for CUDA backend.')
+
+    _set_env_before_cuda_init('CUBLAS_WORKSPACE_CONFIG', _CUBLAS_WORKSPACE_CONFIG)
+    _set_env_before_cuda_init('CUBLASLT_WORKSPACE_SIZE', _CUBLASLT_WORKSPACE_SIZE)
+    _set_torch_precision_policy()
+
+    if validate_device:
+        validate_batch_invariant_device()
+
+    _BATCH_INVARIANT_POLICY_ENABLED = True

lmdeploy/pytorch/backends/cuda/graph_runner.py

@@ -75,6 +75,10 @@ def __init__(
         self.model = model
         self.ctx_mgr = model.ctx_mgr
         self.model_config = model_config
+        from .attention import use_fa3
+        from .batch_invariant import get_fa3_decode_num_splits
+        use_fa3_decoding = bool(
+            use_fa3 and not getattr(model_config, 'use_flash_mla', False) and model_config.head_dim <= 256)
 
         self.meta = CudaGraphMeta(
             max_batchs=max_batches,
@@ -88,8 +92,8 @@ def __init__(
             use_mla_fp8_cache=getattr(self.model_config, 'use_mla_fp8_cache', False),
             use_flash_mla=getattr(self.model_config, 'use_flash_mla', False),
             mla_index_topk=getattr(self.model_config, 'mla_index_topk', None),
-            use_fa3_decoding=(model_config.model_paradigm == 'ar_spec'
-                              and not getattr(model_config, 'use_flash_mla', False)),
+            use_fa3_decoding=use_fa3_decoding,
+            fa3_num_splits=get_fa3_decode_num_splits(),
             is_ssm=len(model_config.states_shapes) > 0,
             use_mrope=model_config.use_mrope,
             block_size=model_config.block_size,

lmdeploy/pytorch/backends/cuda/moe/__init__.py

@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .blocked_fp8 import TritonFusedMoEBlockedF8Builder  # noqa: F401
 from .default import TritonFusedMoEBuilder  # noqa: F401
+from .topk import CudaSoftmaxTopKBuilder  # noqa: F401
 from .w8a8 import TritonFusedMoEW8A8Builder  # noqa: F401