From b83dda6c2176dde72caec2ad4f9eca4ab6e2ba84 Mon Sep 17 00:00:00 2001 From: grimoire Date: Tue, 2 Jun 2026 16:01:19 +0800 Subject: [PATCH 1/8] add batch invariant flag --- lmdeploy/cli/serve.py | 2 + lmdeploy/cli/utils.py | 9 ++ lmdeploy/messages.py | 5 ++ lmdeploy/pytorch/backends/base.py | 6 ++ .../pytorch/backends/cuda/batch_invariant.py | 87 +++++++++++++++++++ lmdeploy/pytorch/backends/cuda/op_backend.py | 6 ++ lmdeploy/pytorch/backends/selector.py | 7 ++ lmdeploy/pytorch/config.py | 1 + lmdeploy/pytorch/engine/config_builder.py | 5 ++ .../pytorch/engine/executor/ray_executor.py | 4 +- .../pytorch/engine/executor/uni_executor.py | 3 + 11 files changed, 134 insertions(+), 1 deletion(-) create mode 100644 lmdeploy/pytorch/backends/cuda/batch_invariant.py diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index b5016d801b..0424730157 100644 --- a/lmdeploy/cli/serve.py +++ b/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, diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index 2c9f8034f6..c7d67c2fbf 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -646,6 +646,15 @@ 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 backends.') + @staticmethod def communicator(parser): return parser.add_argument('--communicator', diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index c262bdbe03..b959ace14e 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -366,6 +366,8 @@ 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. 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 +428,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 +476,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}' diff --git a/lmdeploy/pytorch/backends/base.py b/lmdeploy/pytorch/backends/base.py index 590727ff53..1eee52ca66 100644 --- a/lmdeploy/pytorch/backends/base.py +++ b/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.') diff --git a/lmdeploy/pytorch/backends/cuda/batch_invariant.py b/lmdeploy/pytorch/backends/cuda/batch_invariant.py new file mode 100644 index 0000000000..f055061a13 --- /dev/null +++ b/lmdeploy/pytorch/backends/cuda/batch_invariant.py @@ -0,0 +1,87 @@ +# 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' + + +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') + + matmul_backend = torch.backends.cuda.matmul + if hasattr(matmul_backend, 'allow_fp16_reduced_precision_reduction'): + matmul_backend.allow_fp16_reduced_precision_reduction = False + if hasattr(matmul_backend, 'allow_bf16_reduced_precision_reduction'): + matmul_backend.allow_bf16_reduced_precision_reduction = False + + +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. + """ + 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() diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py index 5608bac35f..8c4a182c12 100644 --- a/lmdeploy/pytorch/backends/cuda/op_backend.py +++ b/lmdeploy/pytorch/backends/cuda/op_backend.py @@ -282,3 +282,9 @@ def device_count(): def support_ray(): """Support ray.""" return True + + @staticmethod + def apply_backend_policy(backend_config: BackendConfig, validate_device: bool = False): + """Apply CUDA backend-specific runtime policy.""" + from .batch_invariant import apply_batch_invariant_policy + apply_batch_invariant_policy(backend_config, validate_device=validate_device) diff --git a/lmdeploy/pytorch/backends/selector.py b/lmdeploy/pytorch/backends/selector.py index 1164cbe693..c17b108985 100644 --- a/lmdeploy/pytorch/backends/selector.py +++ b/lmdeploy/pytorch/backends/selector.py @@ -1,4 +1,5 @@ # Copyright (c) OpenMMLab. All rights reserved. +from lmdeploy.pytorch.config import BackendConfig from lmdeploy.pytorch.devices import DeviceContext, get_device_manager @@ -40,3 +41,9 @@ def init_backend(backend_type: str): """Init device backend.""" backend = get_backend(backend_type) backend.init() + + +def apply_backend_policy(backend_type: str, backend_config: BackendConfig, validate_device: bool = False): + """Apply backend-specific runtime policy.""" + backend = get_backend(backend_type) + backend.apply_backend_policy(backend_config, validate_device=validate_device) diff --git a/lmdeploy/pytorch/config.py b/lmdeploy/pytorch/config.py index af26839855..c3d29331f8 100644 --- a/lmdeploy/pytorch/config.py +++ b/lmdeploy/pytorch/config.py @@ -72,6 +72,7 @@ class BackendConfig: """Backend config.""" eager_mode: bool = True device_type: str = 'cuda' + enable_batch_invariant: bool = False @dataclass diff --git a/lmdeploy/pytorch/engine/config_builder.py b/lmdeploy/pytorch/engine/config_builder.py index f80132eb92..cce1c730eb 100644 --- a/lmdeploy/pytorch/engine/config_builder.py +++ b/lmdeploy/pytorch/engine/config_builder.py @@ -3,6 +3,7 @@ import os from lmdeploy.messages import PytorchEngineConfig, SpeculativeConfig +from lmdeploy.pytorch.backends.selector import apply_backend_policy from lmdeploy.pytorch.config import ( BackendConfig, CacheConfig, @@ -27,6 +28,9 @@ def update_engine_config(engine_config: PytorchEngineConfig): else: engine_config = copy.deepcopy(engine_config) + if engine_config.enable_batch_invariant: + backend_config = ConfigBuilder.build_backend_config(engine_config) + apply_backend_policy(engine_config.device_type, backend_config) if engine_config.max_batch_size is None: engine_config.max_batch_size = get_max_batch_size(engine_config.device_type) @@ -82,6 +86,7 @@ def build_backend_config(engine_config: PytorchEngineConfig): backend_config = BackendConfig( eager_mode=engine_config.eager_mode, device_type=engine_config.device_type, + enable_batch_invariant=engine_config.enable_batch_invariant, ) return backend_config diff --git a/lmdeploy/pytorch/engine/executor/ray_executor.py b/lmdeploy/pytorch/engine/executor/ray_executor.py index 5f6442b514..f36c5d6ed1 100644 --- a/lmdeploy/pytorch/engine/executor/ray_executor.py +++ b/lmdeploy/pytorch/engine/executor/ray_executor.py @@ -12,7 +12,7 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from lmdeploy.pytorch import envs as _envs -from lmdeploy.pytorch.backends.selector import init_backend +from lmdeploy.pytorch.backends.selector import apply_backend_policy, init_backend from lmdeploy.pytorch.config import BackendConfig, CacheConfig, DistConfig, MiscConfig, ModelConfig, SpecDecodeConfig from lmdeploy.pytorch.devices import DeviceContext, get_device_manager from lmdeploy.pytorch.disagg.conn.protocol import DistServeInitRequest, DistServeKVTransferEndpointInfo @@ -170,6 +170,8 @@ def __init__( specdecode_config: SpecDecodeConfig = None, trust_remote_code: bool = False ): + if backend_config.enable_batch_invariant: + apply_backend_policy(device_type, backend_config, validate_device=True) init_backend(device_type) try_import_deeplink(device_type) diff --git a/lmdeploy/pytorch/engine/executor/uni_executor.py b/lmdeploy/pytorch/engine/executor/uni_executor.py index 6bf8389c07..8e9c1cd372 100644 --- a/lmdeploy/pytorch/engine/executor/uni_executor.py +++ b/lmdeploy/pytorch/engine/executor/uni_executor.py @@ -1,6 +1,7 @@ # Copyright (c) OpenMMLab. All rights reserved. import asyncio +from lmdeploy.pytorch.backends.selector import apply_backend_policy from lmdeploy.pytorch.config import BackendConfig, CacheConfig, DistConfig, MiscConfig, ModelConfig, SpecDecodeConfig from lmdeploy.pytorch.devices import DeviceContext from lmdeploy.pytorch.disagg.conn.protocol import DistServeInitRequest, DistServeKVTransferEndpointInfo @@ -40,6 +41,8 @@ def __init__( specdecode_config=specdecode_config, trust_remote_code=trust_remote_code) + if backend_config.enable_batch_invariant: + apply_backend_policy(device_type, backend_config, validate_device=True) self.device_ctx = DeviceContext(device_type=device_type) self.model_agent = build_model_agent( model_path=model_path, From cc350a14e1c412d9cc99be46b2055c7ec6df766d Mon Sep 17 00:00:00 2001 From: grimoire Date: Thu, 4 Jun 2026 12:45:38 +0800 Subject: [PATCH 2/8] topk --- lmdeploy/cli/utils.py | 3 +- lmdeploy/messages.py | 2 + .../pytorch/backends/cuda/batch_invariant.py | 40 ++++++++- .../pytorch/backends/cuda/moe/__init__.py | 1 + lmdeploy/pytorch/backends/cuda/moe/topk.py | 87 +++++++++++++++++++ lmdeploy/pytorch/backends/cuda/op_backend.py | 3 + lmdeploy/pytorch/engine/config_builder.py | 6 ++ lmdeploy/pytorch/engine/model_agent/agent.py | 1 + lmdeploy/pytorch/envs.py | 3 + lmdeploy/pytorch/model_inputs.py | 1 + lmdeploy/pytorch/models/qwen3_5_moe.py | 44 ++++++++-- 11 files changed, 181 insertions(+), 10 deletions(-) create mode 100644 lmdeploy/pytorch/backends/cuda/moe/topk.py diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index c7d67c2fbf..3e103a250d 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -653,7 +653,8 @@ def enable_batch_invariant(parser): return parser.add_argument('--enable-batch-invariant', action='store_true', default=False, - help='Enable batch-invariant greedy inference for supported PyTorch backends.') + help='Enable batch-invariant greedy inference for supported PyTorch backends. ' + 'Can also be enabled with LMDEPLOY_ENABLE_BATCH_INVARIANT=1.') @staticmethod def communicator(parser): diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index b959ace14e..d4119c8aa4 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -368,6 +368,8 @@ class PytorchEngineConfig: 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`. 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 diff --git a/lmdeploy/pytorch/backends/cuda/batch_invariant.py b/lmdeploy/pytorch/backends/cuda/batch_invariant.py index f055061a13..9d97943dbe 100644 --- a/lmdeploy/pytorch/backends/cuda/batch_invariant.py +++ b/lmdeploy/pytorch/backends/cuda/batch_invariant.py @@ -12,6 +12,12 @@ _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 _cuda_initialized() -> bool: @@ -43,11 +49,33 @@ def _set_torch_precision_policy(): 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 - if hasattr(matmul_backend, 'allow_fp16_reduced_precision_reduction'): - matmul_backend.allow_fp16_reduced_precision_reduction = False - if hasattr(matmul_backend, 'allow_bf16_reduced_precision_reduction'): - matmul_backend.allow_bf16_reduced_precision_reduction = False + 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): @@ -73,6 +101,8 @@ def apply_batch_invariant_policy(config: BackendConfig, 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 @@ -85,3 +115,5 @@ def apply_batch_invariant_policy(config: BackendConfig, if validate_device: validate_batch_invariant_device() + + _BATCH_INVARIANT_POLICY_ENABLED = True diff --git a/lmdeploy/pytorch/backends/cuda/moe/__init__.py b/lmdeploy/pytorch/backends/cuda/moe/__init__.py index 882a9ec54c..dfe1994fd8 100644 --- a/lmdeploy/pytorch/backends/cuda/moe/__init__.py +++ b/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 diff --git a/lmdeploy/pytorch/backends/cuda/moe/topk.py b/lmdeploy/pytorch/backends/cuda/moe/topk.py new file mode 100644 index 0000000000..afc6f48088 --- /dev/null +++ b/lmdeploy/pytorch/backends/cuda/moe/topk.py @@ -0,0 +1,87 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""CUDA softmax top-k implementations.""" + +import torch + +from lmdeploy.pytorch.backends.default.moe import DefaultSoftmaxTopKImpl +from lmdeploy.pytorch.backends.moe import SoftmaxTopKBuilder + +from ..batch_invariant import is_batch_invariant_policy_enabled + + +def _load_flashinfer(): + try: + import flashinfer + except ImportError as exc: + raise RuntimeError( + 'enable_batch_invariant for CUDA SoftmaxTopK requires flashinfer with stable top_k support. ' + 'Please install flashinfer or disable enable_batch_invariant.') from exc + + if not hasattr(flashinfer, 'top_k'): + raise RuntimeError( + 'enable_batch_invariant for CUDA SoftmaxTopK requires flashinfer.top_k, ' + 'but the installed flashinfer package does not provide it.') + return flashinfer + + +def _canonicalize_topk(topk_weights: torch.Tensor, topk_ids: torch.Tensor): + """Sort selected experts by expert id for canonical combine order.""" + id_order = torch.argsort(topk_ids, dim=-1, stable=True) + topk_weights = topk_weights.gather(-1, id_order) + topk_ids = topk_ids.gather(-1, id_order) + return topk_weights, topk_ids + + +class CudaSoftmaxTopKImpl(DefaultSoftmaxTopKImpl): + """CUDA softmax top-k implementation.""" + + def __init__(self, top_k: int, dim: int = -1, n_groups: int = -1): + super().__init__(top_k, dim, n_groups=n_groups) + self.batch_invariant = is_batch_invariant_policy_enabled() + self.flashinfer = _load_flashinfer() if self.batch_invariant else None + + def _flashinfer_topk(self, x: torch.Tensor, k: int): + topk_weights, topk_ids = self.flashinfer.top_k( + x, + k, + sorted=True, + deterministic=True, + tie_break=1, + ) + return _canonicalize_topk(topk_weights, topk_ids) + + def forward(self, x: torch.Tensor): + """forward.""" + if not self.batch_invariant: + return super().forward(x) + + if self.dim not in (-1, x.dim() - 1): + raise RuntimeError('Batch-invariant CUDA SoftmaxTopK currently supports only the last dimension.') + + routing_weights = torch.softmax(x, dim=-1, dtype=torch.float32) + if self.n_groups > 0: + assert routing_weights.shape[-1] % self.n_groups == 0, ( + f'{routing_weights.shape[-1]} cannot be divided by {self.n_groups}') + per_group_top_k = self.top_k // self.n_groups + group_size = routing_weights.shape[-1] // self.n_groups + group_offsets = self.get_group_offsets(self.n_groups, group_size, routing_weights.device) + grouped_weights = routing_weights.unflatten(-1, (self.n_groups, group_size)) + original_shape = grouped_weights.shape + flat_weights = grouped_weights.reshape(-1, group_size) + topk_weights, topk_ids = self._flashinfer_topk(flat_weights, per_group_top_k) + topk_weights = topk_weights.reshape(*original_shape[:-1], per_group_top_k) + topk_ids = topk_ids.reshape(*original_shape[:-1], per_group_top_k) + topk_ids = (topk_ids + group_offsets).flatten(-2, -1) + topk_weights = topk_weights.flatten(-2, -1) + return _canonicalize_topk(topk_weights, topk_ids) + + return self._flashinfer_topk(routing_weights, self.top_k) + + +class CudaSoftmaxTopKBuilder(SoftmaxTopKBuilder): + """CUDA softmax top-k implementation builder.""" + + @staticmethod + def build(top_k: int, dim: int = -1, n_groups: int = -1): + """build.""" + return CudaSoftmaxTopKImpl(top_k, dim, n_groups=n_groups) diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py index 8c4a182c12..70e3dd3776 100644 --- a/lmdeploy/pytorch/backends/cuda/op_backend.py +++ b/lmdeploy/pytorch/backends/cuda/op_backend.py @@ -49,6 +49,9 @@ def get_layer_impl_builder(cls, layer_type: OpType): elif layer_type == OpType.SiluAndMul: from .activation import TritonSiluAndMulBuilder return TritonSiluAndMulBuilder + elif layer_type == OpType.SoftmaxTopK: + from .moe import CudaSoftmaxTopKBuilder + return CudaSoftmaxTopKBuilder elif layer_type == OpType.LinearW4A16: from .awq_modules import AwqLinearW4A16Builder return AwqLinearW4A16Builder diff --git a/lmdeploy/pytorch/engine/config_builder.py b/lmdeploy/pytorch/engine/config_builder.py index cce1c730eb..652805d7b5 100644 --- a/lmdeploy/pytorch/engine/config_builder.py +++ b/lmdeploy/pytorch/engine/config_builder.py @@ -3,6 +3,7 @@ import os from lmdeploy.messages import PytorchEngineConfig, SpeculativeConfig +from lmdeploy.pytorch import envs as _envs from lmdeploy.pytorch.backends.selector import apply_backend_policy from lmdeploy.pytorch.config import ( BackendConfig, @@ -28,6 +29,11 @@ def update_engine_config(engine_config: PytorchEngineConfig): else: engine_config = copy.deepcopy(engine_config) + if _envs.enable_batch_invariant: + if engine_config.device_type != 'cuda': + raise ValueError('LMDEPLOY_ENABLE_BATCH_INVARIANT is currently supported only for CUDA backend.') + engine_config.enable_batch_invariant = True + if engine_config.enable_batch_invariant: backend_config = ConfigBuilder.build_backend_config(engine_config) apply_backend_policy(engine_config.device_type, backend_config) diff --git a/lmdeploy/pytorch/engine/model_agent/agent.py b/lmdeploy/pytorch/engine/model_agent/agent.py index e6e0cfbde9..ecf99810a6 100644 --- a/lmdeploy/pytorch/engine/model_agent/agent.py +++ b/lmdeploy/pytorch/engine/model_agent/agent.py @@ -1016,6 +1016,7 @@ def _build_model(self): dllm_config=self.misc_config.dllm_config, strategy_factory=self.strategy_factory, enable_return_routed_experts=enable_return_routed_experts, + enable_batch_invariant=self.backend_config.enable_batch_invariant, quant_config=self.model_config.quant_config, fp32_lm_head=self.model_config.fp32_lm_head, tie_word_embeddings=self.model_config.tie_word_embeddings, diff --git a/lmdeploy/pytorch/envs.py b/lmdeploy/pytorch/envs.py index 9ec6dfa3ae..b3b0d866fd 100644 --- a/lmdeploy/pytorch/envs.py +++ b/lmdeploy/pytorch/envs.py @@ -111,6 +111,9 @@ def _patched_get_env( # executor executor_backend = os.getenv('LMDEPLOY_EXECUTOR_BACKEND', None) + # backend policy + enable_batch_invariant = env_to_bool('LMDEPLOY_ENABLE_BATCH_INVARIANT', False) + # torch profiler torch_profile_cpu = env_to_bool('LMDEPLOY_PROFILE_CPU', False) torch_profile_cuda = env_to_bool('LMDEPLOY_PROFILE_CUDA', False) diff --git a/lmdeploy/pytorch/model_inputs.py b/lmdeploy/pytorch/model_inputs.py index 590683431b..6487b45a76 100644 --- a/lmdeploy/pytorch/model_inputs.py +++ b/lmdeploy/pytorch/model_inputs.py @@ -440,6 +440,7 @@ class BuildModelContext: dllm_config: DLLMConfig = None strategy_factory: 'StrategyFactoryBase' = None enable_return_routed_experts: bool = False + enable_batch_invariant: bool = False quant_config: QuantizationConfig = field(default_factory=QuantizationConfig) fp32_lm_head: bool = False tie_word_embeddings: bool = False diff --git a/lmdeploy/pytorch/models/qwen3_5_moe.py b/lmdeploy/pytorch/models/qwen3_5_moe.py index 92bc7f75f7..230ec918d3 100644 --- a/lmdeploy/pytorch/models/qwen3_5_moe.py +++ b/lmdeploy/pytorch/models/qwen3_5_moe.py @@ -10,7 +10,7 @@ from lmdeploy.pytorch.distributed import get_dist_manager from lmdeploy.pytorch.model_inputs import StepContextManager from lmdeploy.pytorch.nn import RMSNorm -from lmdeploy.pytorch.nn.moe import build_fused_moe +from lmdeploy.pytorch.nn.moe import SoftmaxTopK, build_fused_moe from lmdeploy.pytorch.weight_loader.model_weight_loader import load_weight from .interns1_pro_time_series import InternS1ProTimeSeriesModel @@ -28,6 +28,8 @@ from .qwen3_5 import Qwen3_5VisionModel as Qwen3_5MoeVisionModel from .qwen3_vl import Qwen3VLInputProcessor as Qwen3_5MoeInputProcessor +_BATCH_INVARIANT_SHARED_EXPERT_GATE_OUT_FEATURES = 8 + class Qwen3_5MoeTopKRouter(nn.Module): @@ -37,18 +39,50 @@ def __init__(self, config, dtype: torch.dtype | None = None, device: torch.devic self.num_experts = config.num_experts self.hidden_dim = config.hidden_size self.weight = nn.Parameter(torch.zeros(self.num_experts, self.hidden_dim, dtype=dtype, device=device)) + self.softmax_topk = SoftmaxTopK(self.top_k) def forward(self, hidden_states): hidden_states = hidden_states.reshape(-1, self.hidden_dim) router_logits = F.linear(hidden_states, self.weight) # (seq_len, num_experts) - router_logits = torch.nn.functional.softmax(router_logits, dtype=torch.float, dim=-1) - router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) # (seq_len, top_k) + router_top_value, router_indices = self.softmax_topk(router_logits) router_top_value /= router_top_value.sum(dim=-1, keepdim=True) router_top_value = router_top_value.to(router_logits.dtype) router_scores = router_top_value return router_logits, router_scores, router_indices +class Qwen3_5MoeSharedExpertGate(nn.Module): + """Shared expert gate with a padded batch-invariant CUDA path.""" + + def __init__(self, + hidden_size: int, + dtype: torch.dtype | None = None, + device: torch.device | None = None): + super().__init__() + self.enable_batch_invariant = get_build_model_context().enable_batch_invariant + out_features = 1 + if self.enable_batch_invariant: + out_features = _BATCH_INVARIANT_SHARED_EXPERT_GATE_OUT_FEATURES + self.weight = nn.Parameter(torch.empty(out_features, hidden_size, dtype=dtype, device=device)) + if self.enable_batch_invariant: + self.weight.weight_loader = self.weight_loader + + def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor): + """Load the scalar checkpoint gate into row 0 of the padded gate.""" + assert loaded_weight.size(0) == 1 + assert loaded_weight.size(1) == param.size(1) + assert param.size(0) >= loaded_weight.size(0) + param.data.zero_() + param.data[:loaded_weight.size(0)].copy_(loaded_weight) + + def forward(self, hidden_states: torch.Tensor): + """Forward.""" + gate = F.linear(hidden_states, self.weight) + if self.enable_batch_invariant: + gate = gate[..., :1] + return gate.sigmoid() + + class Qwen3_5MoeSparseMoeBlock(nn.Module): """Sparse MoE block.""" @@ -94,7 +128,7 @@ def __init__(self, all_reduce=False, prefix=add_prefix('shared_expert', prefix), ) - self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False, device=device, dtype=dtype) + self.shared_expert_gate = Qwen3_5MoeSharedExpertGate(config.hidden_size, device=device, dtype=dtype) # get all reduce dist_ctx = get_dist_manager().current_context() @@ -119,7 +153,7 @@ def forward(self, hidden_states: torch.Tensor, all_routed_experts: torch.Tensor ) shared_states = self.shared_expert(hidden_states) - shared_states = self.shared_expert_gate(hidden_states).sigmoid() * shared_states + shared_states = self.shared_expert_gate(hidden_states) * shared_states out_states += shared_states out_states = out_states.reshape(batch_size, sequence_length, -1) From 33094b381122954fe2089e444c931d886849a0a5 Mon Sep 17 00:00:00 2001 From: grimoire Date: Thu, 4 Jun 2026 16:56:10 +0800 Subject: [PATCH 3/8] update fa3 --- .../backends/cuda/attention/__init__.py | 24 +++++++ .../pytorch/backends/cuda/attention/fa3.py | 72 +++++++++++++++---- .../pytorch/backends/cuda/batch_invariant.py | 5 ++ .../pytorch/backends/cuda/graph_runner.py | 8 ++- lmdeploy/pytorch/backends/cuda/op_backend.py | 27 ++++--- .../pytorch/kernels/cuda/pagedattention.py | 12 ++++ lmdeploy/pytorch/models/utils/cudagraph.py | 10 ++- 7 files changed, 129 insertions(+), 29 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/attention/__init__.py b/lmdeploy/pytorch/backends/cuda/attention/__init__.py index 37251f7b80..d38ffc3d6a 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/__init__.py +++ b/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 diff --git a/lmdeploy/pytorch/backends/cuda/attention/fa3.py b/lmdeploy/pytorch/backends/cuda/attention/fa3.py index 576ddb2dc5..5c92cd1fcf 100644 --- a/lmdeploy/pytorch/backends/cuda/attention/fa3.py +++ b/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,7 +149,7 @@ 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, @@ -134,6 +157,7 @@ def _decoding_speculative( 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, diff --git a/lmdeploy/pytorch/backends/cuda/batch_invariant.py b/lmdeploy/pytorch/backends/cuda/batch_invariant.py index 9d97943dbe..150d51cadf 100644 --- a/lmdeploy/pytorch/backends/cuda/batch_invariant.py +++ b/lmdeploy/pytorch/backends/cuda/batch_invariant.py @@ -20,6 +20,11 @@ def is_batch_invariant_policy_enabled() -> bool: 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: diff --git a/lmdeploy/pytorch/backends/cuda/graph_runner.py b/lmdeploy/pytorch/backends/cuda/graph_runner.py index 601f2c0dce..08db379dea 100644 --- a/lmdeploy/pytorch/backends/cuda/graph_runner.py +++ b/lmdeploy/pytorch/backends/cuda/graph_runner.py @@ -74,6 +74,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, @@ -87,8 +91,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, diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py index 70e3dd3776..89070ecce4 100644 --- a/lmdeploy/pytorch/backends/cuda/op_backend.py +++ b/lmdeploy/pytorch/backends/cuda/op_backend.py @@ -143,9 +143,11 @@ def update_meta_flashattn(cls, attn_metadata, step_context): from lmdeploy.pytorch.models.utils.cudagraph import _get_meta_flashattn from .attention import _normalize_sliding_window + from .batch_invariant import get_fa3_decode_num_splits batch_size = attn_metadata.q_seqlens.size(0) max_seqlen_q = step_context.input_ids.size(1) // batch_size window_size = _normalize_sliding_window(step_context.model_config.sliding_window) + num_splits = get_fa3_decode_num_splits() # FA3's mha_fwd internally computes seqlen_k as # page_table.size(1) * page_size when cu_seqlens_k is None # (paged KV without varlen_k). get_scheduler_metadata must use @@ -163,11 +165,19 @@ def update_meta_flashattn(cls, attn_metadata, step_context): qkv_dtype=step_context.model_config.dtype, page_size=step_context.model_config.block_size, window_size=window_size, + num_splits=num_splits, ) attn_metadata.scheduler_metadata = scheduler_metadata attn_metadata.max_kv_seqlen = max_seqlen_k return attn_metadata + @staticmethod + def _should_build_fa3_metadata(model_config: ModelConfig): + """Return whether FA3 metadata should be built.""" + from .attention import use_fa3 + + return bool(use_fa3 and not model_config.use_flash_mla and model_config.head_dim <= 256) + @classmethod def update_chunked_gated_delta_rule_meta(cls, attn_metadata, step_context): try: @@ -207,7 +217,8 @@ def update_step_context(cls, step_context): kv_start_loc = None kv_flatten_size = None use_flash_mla = step_context.model_config.use_flash_mla - use_flash_attn3_decoding = step_context.model_config.model_paradigm == 'ar_spec' + use_flash_attn3_decoding = cls._should_build_fa3_metadata(step_context.model_config) + requires_flash_attn3_decoding = step_context.model_config.model_paradigm == 'ar_spec' and not use_flash_mla # pad and cumsum requires 4 kernels, so we fuse seqlens cumsum into one kernel seqlens = torch.stack([q_seqlens, kv_seqlens], dim=0) @@ -238,17 +249,11 @@ def update_step_context(cls, step_context): decode_query_len = step_context.input_ids.size(1) // q_seqlens.size(0) cls.update_meta_flashmla(attn_metadata, model_config, decode_query_len) elif use_flash_attn3_decoding: - from .attention import use_fa3 - if not use_fa3: - sm = torch.cuda.get_device_capability() - cuda_ver = torch.version.cuda or 'N/A' - raise RuntimeError( - f'Speculative decoding on CUDA requires FlashAttention-3 (FA3), ' - f'which needs SM80+ (Ampere and above) with CUDA >= 12.3 and ' - f'flash-attn installed. Detected: SM{sm[0]}.{sm[1]}, CUDA {cuda_ver}. ' - f'Please ensure your GPU meets SM80+, CUDA >= 12.3, and flash-attn ' - f'is installed, or disable speculative decoding.') attn_metadata = cls.update_meta_flashattn(attn_metadata, step_context) + elif requires_flash_attn3_decoding: + raise RuntimeError( + 'Speculative decoding on CUDA requires FA3 attention. Please ensure FA3 is available and the ' + 'attention configuration supports FA3, or disable speculative decoding.') # update chunk gated delta indices is_gated_delta = step_context.model_config.is_gated_delta diff --git a/lmdeploy/pytorch/kernels/cuda/pagedattention.py b/lmdeploy/pytorch/kernels/cuda/pagedattention.py index 3f4955bb61..8c8a680e19 100644 --- a/lmdeploy/pytorch/kernels/cuda/pagedattention.py +++ b/lmdeploy/pytorch/kernels/cuda/pagedattention.py @@ -672,6 +672,17 @@ def _get_split_k(device_idx: int, head_grid: int, batch_size: int, num_warps: in return SPLIT_K +def _cap_split_k_by_page_table(split_k: int, page_table: Tensor): + """Cap split-K by the number of available KV pages. + + Empty split partitions do no useful work but still allocate accumulator memory and participate in the reduction. + Keep the cap power-of-two so the Triton reduction shape remains valid. + """ + max_pages = max(int(page_table.size(1)), 1) + max_useful_split = max(4, triton.next_power_of_2(max_pages)) + return min(split_k, max_useful_split) + + @triton.jit def _bar_sync(): """CTA-internal barrier (__syncthreads equivalent via PTX bar.sync 0).""" @@ -892,6 +903,7 @@ def _get_block_d(Lk): is_fp8_scalar = quant_policy in (QuantPolicy.FP8, QuantPolicy.FP8_E5M2) SPLIT_K = _get_split_k(q.device.index, grid_1, batch, num_warps) + SPLIT_K = _cap_split_k_by_page_table(SPLIT_K, page_table) if quant_policy == QuantPolicy.INT4 or quant_policy == QuantPolicy.TURBO_QUANT: acc = q.new_empty(num_tokens, head, SPLIT_K, o.shape[-1] + 2, dtype=torch.float32) diff --git a/lmdeploy/pytorch/models/utils/cudagraph.py b/lmdeploy/pytorch/models/utils/cudagraph.py index bbcf933585..55335ff1bc 100644 --- a/lmdeploy/pytorch/models/utils/cudagraph.py +++ b/lmdeploy/pytorch/models/utils/cudagraph.py @@ -79,6 +79,7 @@ class CudaGraphMeta: use_flash_mla: bool = False mla_index_topk: int | None = None use_fa3_decoding: bool = False + fa3_num_splits: int = 0 is_ssm: bool = False use_mrope: bool = False block_size: int = 64 @@ -115,7 +116,7 @@ def make_output_buffers(self, output): return output_buffers def update_meta_flashattn(self, batch_size: int, max_seqlen_q: int, block_size: int, max_seqlen_k: int, - cache_seqlens: torch.Tensor): + cache_seqlens: torch.Tensor, num_splits: int = 0): """Update meta flashattn.""" ctx_mgr = get_step_ctx_manager() step_ctx = ctx_mgr.current_context() @@ -143,6 +144,7 @@ def update_meta_flashattn(self, batch_size: int, max_seqlen_q: int, block_size: qkv_dtype=torch_dtype, page_size=block_size, window_size=window_size, + num_splits=num_splits, ) return scheduler_metadata @@ -194,14 +196,15 @@ def make_buffers_cudagraph(self, graph_meta: CudaGraphMeta, input_ids: Tensor, p is_fp8_kvcache=graph_meta.use_mla_fp8_cache, topk=index_topk) - # use fa3 decode kernel for spec decode + # use FA3 decode kernel elif graph_meta.use_fa3_decoding is True: max_seqlen_k = graph_meta.num_blocks * graph_meta.block_size input_buffers['scheduler_metadata'] = self.update_meta_flashattn(graph_meta.max_batchs, decode_query_len, block_size=graph_meta.block_size, max_seqlen_k=max_seqlen_k, - cache_seqlens=input_buffers['kv_seqlens']) + cache_seqlens=input_buffers['kv_seqlens'], + num_splits=graph_meta.fa3_num_splits) # mrope if graph_meta.use_mrope: @@ -295,6 +298,7 @@ def fill_buffers_cudagraph(self, graph_meta: CudaGraphMeta, input_ids: Tensor, p block_size=graph_meta.block_size, max_seqlen_k=max_seqlen_k, cache_seqlens=input_buffers['kv_seqlens'], + num_splits=graph_meta.fa3_num_splits, ) num_meta = scheduler_metadata.size(0) input_buffers['scheduler_metadata'][:num_meta] = scheduler_metadata From 6af97c9f3b9daf4f96e3d3e2d4ede1e6a10e666b Mon Sep 17 00:00:00 2001 From: grimoire Date: Fri, 5 Jun 2026 14:25:07 +0800 Subject: [PATCH 4/8] fix fa3 --- lmdeploy/pytorch/backends/cuda/op_backend.py | 1 + lmdeploy/pytorch/models/utils/cudagraph.py | 23 +++++++++++++------- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py index 89070ecce4..a759403837 100644 --- a/lmdeploy/pytorch/backends/cuda/op_backend.py +++ b/lmdeploy/pytorch/backends/cuda/op_backend.py @@ -121,6 +121,7 @@ def get_v_block_shape( def update_meta_flashmla(cls, attn_metadata, model_config: ModelConfig, decoding_query_len: int): """Update meta for flashmla.""" import flash_mla + num_attention_heads, _ = model_config.get_num_qkv_head_by_tp() num_attention_heads *= decoding_query_len is_fp8_kvcache = model_config.use_mla_fp8_cache diff --git a/lmdeploy/pytorch/models/utils/cudagraph.py b/lmdeploy/pytorch/models/utils/cudagraph.py index 55335ff1bc..defb1d3439 100644 --- a/lmdeploy/pytorch/models/utils/cudagraph.py +++ b/lmdeploy/pytorch/models/utils/cudagraph.py @@ -115,8 +115,13 @@ def make_output_buffers(self, output): output_buffers = output return output_buffers - def update_meta_flashattn(self, batch_size: int, max_seqlen_q: int, block_size: int, max_seqlen_k: int, - cache_seqlens: torch.Tensor, num_splits: int = 0): + def update_meta_flashattn(self, + batch_size: int, + max_seqlen_q: int, + block_size: int, + max_seqlen_k: int, + cache_seqlens: torch.Tensor, + num_splits: int = 0): """Update meta flashattn.""" ctx_mgr = get_step_ctx_manager() step_ctx = ctx_mgr.current_context() @@ -199,12 +204,14 @@ def make_buffers_cudagraph(self, graph_meta: CudaGraphMeta, input_ids: Tensor, p # use FA3 decode kernel elif graph_meta.use_fa3_decoding is True: max_seqlen_k = graph_meta.num_blocks * graph_meta.block_size - input_buffers['scheduler_metadata'] = self.update_meta_flashattn(graph_meta.max_batchs, - decode_query_len, - block_size=graph_meta.block_size, - max_seqlen_k=max_seqlen_k, - cache_seqlens=input_buffers['kv_seqlens'], - num_splits=graph_meta.fa3_num_splits) + input_buffers['scheduler_metadata'] = self.update_meta_flashattn( + graph_meta.max_batchs, + decode_query_len, + block_size=graph_meta.block_size, + max_seqlen_k=max_seqlen_k, + cache_seqlens=input_buffers['kv_seqlens'], + num_splits=graph_meta.fa3_num_splits, + ) # mrope if graph_meta.use_mrope: From eacaa9d9ee3c9358d02dd25f1f957eccc740e6ee Mon Sep 17 00:00:00 2001 From: grimoire Date: Tue, 9 Jun 2026 14:09:55 +0800 Subject: [PATCH 5/8] fix mtp --- lmdeploy/pytorch/backends/cuda/graph_runner.py | 3 ++- lmdeploy/pytorch/strategies/ar_spec/cudagraph.py | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/graph_runner.py b/lmdeploy/pytorch/backends/cuda/graph_runner.py index 08db379dea..54f179447e 100644 --- a/lmdeploy/pytorch/backends/cuda/graph_runner.py +++ b/lmdeploy/pytorch/backends/cuda/graph_runner.py @@ -230,6 +230,7 @@ def get_graph_key(self, input_ids: torch.Tensor, position_ids: torch.Tensor, pas context = self.ctx_mgr.current_context() is_decoding = context.global_is_decoding() batch_size = attn_metadata.q_seqlens.size(0) + query_len = input_ids.size(1) // batch_size meta = self.get_meta() enable_microbatch = get_step_ctx_manager().current_context().enable_microbatch # for draft model to distinguish inputs from target model and itself @@ -240,7 +241,7 @@ def get_graph_key(self, input_ids: torch.Tensor, position_ids: torch.Tensor, pas batch_size = self._get_capture_tokens(batch_size) else: batch_size = self._get_capture_tokens(meta.padding_batch_size) - return (batch_size, is_decoding, enable_microbatch, target_hidden_size) + return (batch_size, is_decoding, enable_microbatch, target_hidden_size, query_len) def _prepare_inputs(self, **kwargs): """Prepare inputs.""" diff --git a/lmdeploy/pytorch/strategies/ar_spec/cudagraph.py b/lmdeploy/pytorch/strategies/ar_spec/cudagraph.py index e77fbb715d..f851b6e1f9 100644 --- a/lmdeploy/pytorch/strategies/ar_spec/cudagraph.py +++ b/lmdeploy/pytorch/strategies/ar_spec/cudagraph.py @@ -12,8 +12,10 @@ def __init__(self, num_spec_tokens: int, method: str): def get_max_tokens(self, batch_size: int, origin_batch_size: int, num_tokens: int) -> int: """Get max tokens.""" - # only eagle3 has two sets of cudagraph due to different target_hidden_size - if num_tokens == origin_batch_size and self.method == 'eagle3': + # Draft speculative decoding has both wide verification forwards and + # single-token iterative forwards. Keep their graph buffers shaped like + # the runtime query layout even when target_hidden_size is identical. + if num_tokens == origin_batch_size: return batch_size return batch_size * (self.num_spec_tokens + 1) From 66804efb4fb18e79faa1ccec1f352d4ad94303b5 Mon Sep 17 00:00:00 2001 From: grimoire Date: Tue, 9 Jun 2026 15:47:35 +0800 Subject: [PATCH 6/8] fix gdr --- .../pytorch/kernels/cuda/gated_delta_rule.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/lmdeploy/pytorch/kernels/cuda/gated_delta_rule.py b/lmdeploy/pytorch/kernels/cuda/gated_delta_rule.py index 0f6385f918..6cd69aedd1 100644 --- a/lmdeploy/pytorch/kernels/cuda/gated_delta_rule.py +++ b/lmdeploy/pytorch/kernels/cuda/gated_delta_rule.py @@ -1,10 +1,20 @@ # Copyright (c) OpenMMLab. All rights reserved. +import inspect from collections.abc import Sequence import tilelang import tilelang.language as T import torch +_TL_SHFL_SYNC_VALUE_FIRST = next(iter(inspect.signature(T.shfl_sync).parameters)) == 'value' + + +def _shfl_sync(value, src_lane: int, mask: int = 0xFFFFFFFF): + """TileLang 0.1.8 and 0.1.9+ use different Python shfl_sync signatures.""" + if _TL_SHFL_SYNC_VALUE_FIRST: + return T.shfl_sync(value, src_lane, mask=mask) + return T.shfl_sync(mask, value, src_lane) + @T.macro def load_state_tile_from_smem(h_smem: T.Buffer, h_local: T.Buffer, k_off, v_warp_off, K: int, k_per_thr: int, @@ -368,7 +378,7 @@ def fused_recurrent_gated_delta_rule_default_main( g_exp = T.exp(g) else: g_exp = 0.0 - g_exp = T.shfl_sync(0xFFFFFFFF, g_exp, 0) + g_exp = _shfl_sync(g_exp, 0) else: g_exp = 1.0 if use_beta: @@ -376,7 +386,7 @@ def fused_recurrent_gated_delta_rule_default_main( beta = T.cast(Beta[b_id, seq_id, hv_id], T.float32) else: beta = 0.0 - beta = T.shfl_sync(0xFFFFFFFF, beta, 0) + beta = _shfl_sync(beta, 0) else: beta = 1.0 @@ -495,7 +505,7 @@ def fused_recurrent_gated_delta_rule_transposed_main( g_exp = T.exp(g) else: g_exp = 0.0 - g_exp = T.shfl_sync(0xFFFFFFFF, g_exp, 0) + g_exp = _shfl_sync(g_exp, 0) else: g_exp = 1.0 if use_beta: @@ -503,7 +513,7 @@ def fused_recurrent_gated_delta_rule_transposed_main( beta = T.cast(Beta[b_id, seq_id, hv_id], T.float32) else: beta = 0.0 - beta = T.shfl_sync(0xFFFFFFFF, beta, 0) + beta = _shfl_sync(beta, 0) else: beta = 1.0 From 63e2e4caa299080bf7b724fd8485dbdcd4d90448 Mon Sep 17 00:00:00 2001 From: grimoire Date: Wed, 10 Jun 2026 14:16:02 +0800 Subject: [PATCH 7/8] update check --- lmdeploy/cli/utils.py | 6 ++++-- lmdeploy/messages.py | 4 +++- lmdeploy/pytorch/engine/config_builder.py | 8 ++++++++ 3 files changed, 15 insertions(+), 3 deletions(-) diff --git a/lmdeploy/cli/utils.py b/lmdeploy/cli/utils.py index 3e103a250d..ca802f1626 100644 --- a/lmdeploy/cli/utils.py +++ b/lmdeploy/cli/utils.py @@ -653,8 +653,10 @@ def enable_batch_invariant(parser): return parser.add_argument('--enable-batch-invariant', action='store_true', default=False, - help='Enable batch-invariant greedy inference for supported PyTorch backends. ' - 'Can also be enabled with LMDEPLOY_ENABLE_BATCH_INVARIANT=1.') + 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): diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index d4119c8aa4..976c8e02a7 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -369,7 +369,9 @@ class PytorchEngineConfig: 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`. + `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 diff --git a/lmdeploy/pytorch/engine/config_builder.py b/lmdeploy/pytorch/engine/config_builder.py index 652805d7b5..8e2c488eae 100644 --- a/lmdeploy/pytorch/engine/config_builder.py +++ b/lmdeploy/pytorch/engine/config_builder.py @@ -35,6 +35,14 @@ def update_engine_config(engine_config: PytorchEngineConfig): engine_config.enable_batch_invariant = True if engine_config.enable_batch_invariant: + unsupported = [] + if engine_config.ep > 1: + unsupported.append('expert parallelism (ep > 1)') + if engine_config.enable_eplb: + unsupported.append('EPLB') + if unsupported: + unsupported_features = ', '.join(unsupported) + raise ValueError(f'enable_batch_invariant currently does not support {unsupported_features}.') backend_config = ConfigBuilder.build_backend_config(engine_config) apply_backend_policy(engine_config.device_type, backend_config) if engine_config.max_batch_size is None: From 29666fe9419f881bc2e17b3e38b44776ea7ff75e Mon Sep 17 00:00:00 2001 From: grimoire Date: Wed, 10 Jun 2026 15:15:03 +0800 Subject: [PATCH 8/8] fix copilot --- lmdeploy/pytorch/backends/cuda/batch_invariant.py | 5 +++-- lmdeploy/pytorch/backends/cuda/op_backend.py | 3 ++- lmdeploy/pytorch/kernels/cuda/pagedattention.py | 2 +- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/lmdeploy/pytorch/backends/cuda/batch_invariant.py b/lmdeploy/pytorch/backends/cuda/batch_invariant.py index 150d51cadf..5914a7c920 100644 --- a/lmdeploy/pytorch/backends/cuda/batch_invariant.py +++ b/lmdeploy/pytorch/backends/cuda/batch_invariant.py @@ -38,9 +38,10 @@ def _set_env_before_cuda_init(name: str, value: str): current = os.environ.get(name) if current == value: return - if current is not None and _cuda_initialized(): + if _cuda_initialized(): + found = current if current is not None else '' raise RuntimeError( - f'enable_batch_invariant requires {name}={value}, but found {name}={current} ' + f'enable_batch_invariant requires {name}={value}, but found {name}={found} ' '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: diff --git a/lmdeploy/pytorch/backends/cuda/op_backend.py b/lmdeploy/pytorch/backends/cuda/op_backend.py index a759403837..f728f83a57 100644 --- a/lmdeploy/pytorch/backends/cuda/op_backend.py +++ b/lmdeploy/pytorch/backends/cuda/op_backend.py @@ -253,7 +253,8 @@ def update_step_context(cls, step_context): attn_metadata = cls.update_meta_flashattn(attn_metadata, step_context) elif requires_flash_attn3_decoding: raise RuntimeError( - 'Speculative decoding on CUDA requires FA3 attention. Please ensure FA3 is available and the ' + 'Speculative decoding on CUDA requires FA3 attention. Please ensure FlashAttention-3 is ' + 'installed, the GPU supports SM80+, CUDA >= 12.3 is available, head_dim <= 256, and the ' 'attention configuration supports FA3, or disable speculative decoding.') # update chunk gated delta indices diff --git a/lmdeploy/pytorch/kernels/cuda/pagedattention.py b/lmdeploy/pytorch/kernels/cuda/pagedattention.py index 8c8a680e19..20c0d2f9b1 100644 --- a/lmdeploy/pytorch/kernels/cuda/pagedattention.py +++ b/lmdeploy/pytorch/kernels/cuda/pagedattention.py @@ -679,7 +679,7 @@ def _cap_split_k_by_page_table(split_k: int, page_table: Tensor): Keep the cap power-of-two so the Triton reduction shape remains valid. """ max_pages = max(int(page_table.size(1)), 1) - max_useful_split = max(4, triton.next_power_of_2(max_pages)) + max_useful_split = 1 << (max_pages.bit_length() - 1) return min(split_k, max_useful_split)