diff --git a/custom_ops/gpu_ops/append_attn/gqa_rope_write_cache.cu b/custom_ops/gpu_ops/append_attn/gqa_rope_write_cache.cu index c86ec27dca8..09ed4c18e6d 100644 --- a/custom_ops/gpu_ops/append_attn/gqa_rope_write_cache.cu +++ b/custom_ops/gpu_ops/append_attn/gqa_rope_write_cache.cu @@ -281,6 +281,7 @@ __global__ void GQAVariableLengthNeoxPartialRotarySplitKernel( const int token_idx = linear_index / offset; // token id(第几个token,不分qkv) const int ori_bi = batch_id_per_token[token_idx]; // 第几个batch + if (ori_bi == -1) continue; int cache_kv_len = seq_lens_decoder[ori_bi]; // 这里其实是不需要处理的，但是由于FA3的bug，所以必须！ diff --git a/custom_ops/gpu_ops/append_attn/qwen3_rope.h b/custom_ops/gpu_ops/append_attn/qwen3_rope.h index 42017e42151..ab92ccd38d1 100644 --- a/custom_ops/gpu_ops/append_attn/qwen3_rope.h +++ b/custom_ops/gpu_ops/append_attn/qwen3_rope.h @@ -40,6 +40,7 @@ __global__ void GQAVariableLengthRotarySplitKernel_Qwen3( const int token_idx = linear_index / offset; const int ori_bi = batch_id_per_token[token_idx]; // 第几个batch + if (ori_bi == -1) continue; // padded token for static/CUDA-graph shape int cache_kv_len = seq_lens_decoder[ori_bi]; // 这里其实是不需要处理的，但是由于FA3的bug，所以必须！ diff --git a/custom_ops/gpu_ops/grouped_topk_kernels.cu b/custom_ops/gpu_ops/grouped_topk_kernels.cu index ef5ed8533f0..c0713fee954 100644 --- a/custom_ops/gpu_ops/grouped_topk_kernels.cu +++ b/custom_ops/gpu_ops/grouped_topk_kernels.cu @@ -767,6 +767,8 @@ std::vector GroupedTopkInferDtype( std::vector> GroupedTopkInferShape( const std::vector& gating_output_shape, const std::vector&, + const int n_group, + const int topk_group, const int topk) { auto num_tokens = gating_output_shape[0]; auto num_experts = gating_output_shape[1]; diff --git a/custom_ops/gpu_ops/moe/tritonmoe_preprocess.cu b/custom_ops/gpu_ops/moe/tritonmoe_preprocess.cu index eb680ea744e..5af1de6af69 100644 --- a/custom_ops/gpu_ops/moe/tritonmoe_preprocess.cu +++ b/custom_ops/gpu_ops/moe/tritonmoe_preprocess.cu @@ -29,10 +29,25 @@ std::vector> tritonmoe_preprocessInferShape( const std::vector& topk_ids, int64_t num_experts, int64_t GEMM_BLOCK_SIZE_M) { + // Compute numel from shape; any dim == -1 means dynamic/unknown. + // If the shape is fully static and positive, compute the exact upper bound. + // Otherwise fall back to symbolic (-1) so that SOT treats them as + // SymbolicInt rather than raising "negative dimension" errors. + bool has_dynamic_dim = false; int topk_ids_numel = 1; for (int64_t dim : topk_ids) { + if (dim < 0) { + has_dynamic_dim = true; + break; + } topk_ids_numel *= static_cast(dim); } + + if (has_dynamic_dim) { + // Return symbolic shapes so downstream SOT infer_meta stays valid. + return {{-1}, {-1}, {1}}; + } + int max_num_tokens_padded; if (topk_ids_numel < num_experts + 1) { max_num_tokens_padded = topk_ids_numel * GEMM_BLOCK_SIZE_M; diff --git a/custom_ops/gpu_ops/noaux_tc.cu b/custom_ops/gpu_ops/noaux_tc.cu index b9b2d9d32b7..d0f9c2ebd6c 100644 --- a/custom_ops/gpu_ops/noaux_tc.cu +++ b/custom_ops/gpu_ops/noaux_tc.cu @@ -67,6 +67,8 @@ std::vector NoauxTcInferDtype( std::vector> NoauxTcInferShape( const std::vector& scores_shape, const std::vector&, + const int n_group, + const int topk_group, const int topk) { auto num_tokens = scores_shape[0]; auto topk_values_shape = std::vector{num_tokens, topk}; diff --git a/custom_ops/gpu_ops/noaux_tc_redundant.cu b/custom_ops/gpu_ops/noaux_tc_redundant.cu index d9b8483bce6..f991e620cc6 100644 --- a/custom_ops/gpu_ops/noaux_tc_redundant.cu +++ b/custom_ops/gpu_ops/noaux_tc_redundant.cu @@ -74,6 +74,11 @@ std::vector NoauxTcRedundantInferDtype( std::vector> NoauxTcRedundantInferShape( const std::vector& scores_shape, const std::vector&, + const std::vector&, + const std::vector&, + const std::vector&, + const int n_group, + const int topk_group, const int topk) { auto num_tokens = scores_shape[0]; auto topk_values_shape = std::vector{num_tokens, topk}; diff --git a/fastdeploy/config.py b/fastdeploy/config.py index d959b3ce6d7..e2f15cfd014 100644 --- a/fastdeploy/config.py +++ b/fastdeploy/config.py @@ -1151,6 +1151,15 @@ def init_with_cudagrpah_size( ) self.cudagraph_capture_sizes = dedup_sizes + dedup_prefill_sizes = list(set(self.cudagraph_capture_sizes_prefill)) + if len(dedup_prefill_sizes) < len(self.cudagraph_capture_sizes_prefill): + logger.info( + ("cudagraph prefill sizes specified by model runner" " %s is overridden by config %s"), + self.cudagraph_capture_sizes_prefill, + dedup_prefill_sizes, + ) + self.cudagraph_capture_sizes_prefill = dedup_prefill_sizes + # Sort to make sure cudagraph capture sizes are in descending order self.cudagraph_capture_sizes.sort(reverse=True) self.cudagraph_capture_sizes_prefill.sort(reverse=True) @@ -1214,6 +1223,13 @@ def _set_cudagraph_sizes( # Shape [256, 288, ... 992, 1024] draft_capture_sizes += [32 * i for i in range(9, 33)] + # Shape [1024, 1088, ... 2048] step=64 + draft_capture_sizes += [64 * i for i in range(17, 33)] + # Shape [2048, 2176, ... 4096] step=128 + draft_capture_sizes += [128 * i for i in range(17, 33)] + # Shape [4096, 4352, ... 8192] step=256 + draft_capture_sizes += [256 * i for i in range(17, 33)] + draft_capture_sizes_prefill = draft_capture_sizes.copy() draft_capture_sizes.append(max_capture_size) self.cudagraph_capture_sizes = sorted(draft_capture_sizes) @@ -2283,7 +2299,13 @@ def postprocess(self): # Adjustment GraphOptConfig if self.scheduler_config is not None and self.scheduler_config.splitwise_role == "prefill": - self.graph_opt_config.use_cudagraph = self.graph_opt_config.cudagraph_only_prefill + # Piecewise CUDAGraph for prefill worker: if graph_opt_level >= 1 and not full_cuda_graph, + # reuse the mixed piecewise path (capture_model_prefill_and_mixed) for the prefill worker. + # Otherwise fall back to cudagraph_only_prefill flag (legacy path). + if self.graph_opt_config.graph_opt_level >= 1 and not self.graph_opt_config.full_cuda_graph: + self.graph_opt_config.use_cudagraph = True + else: + self.graph_opt_config.use_cudagraph = self.graph_opt_config.cudagraph_only_prefill if self.load_config is not None and self.load_config.dynamic_load_weight is True: self.graph_opt_config.graph_opt_level = 0 logger.info( diff --git a/fastdeploy/envs.py b/fastdeploy/envs.py index 83859610dae..bb7662a41e5 100644 --- a/fastdeploy/envs.py +++ b/fastdeploy/envs.py @@ -278,22 +278,12 @@ def _validate_split_kv_size(value: int) -> int: # When v1 is enabled, the legacy /clear_load_weight and /update_model_weight # will adopt this new communication pattern. "FD_ENABLE_V1_UPDATE_WEIGHTS": lambda: bool(int(os.getenv("FD_ENABLE_V1_UPDATE_WEIGHTS", "0"))), + # Whether to use GDR CheckpointTransfer for dynamic weight updates. + "FD_USE_GDR_CHECKPOINT_TRANSFER": lambda: bool(int(os.getenv("FD_USE_GDR_CHECKPOINT_TRANSFER", "0"))), # Whether to save the cache of output token for preempted request to storage. "FD_SAVE_OUTPUT_CACHE_FOR_PREEMPTED_REQUEST": lambda: bool( int(os.getenv("FD_SAVE_OUTPUT_CACHE_FOR_PREEMPTED_REQUEST", "1")) ), - # Whether to use GDR CheckpointTransfer for dynamic weight updates. - "FD_USE_GDR_CHECKPOINT_TRANSFER": lambda: bool(int(os.getenv("FD_USE_GDR_CHECKPOINT_TRANSFER", "0"))), - # Whether to enable block-wise CUDA Graph capture/replay. - # When enabled, individual layer forward methods decorated with @block_wise_cuda_graph_wrap - # will be captured and replayed as CUDA Graphs for improved performance. - # Set to 1 to enable; defaults to 0 (disabled). - "FD_USE_BLOCK_WISE_CUDA_GRAPH": lambda: bool(int(os.getenv("FD_USE_BLOCK_WISE_CUDA_GRAPH", "0"))), - # Comma-separated list of token counts to pre-capture for block-wise CUDA Graphs. - # Used during the warmup phase to pre-capture graphs for these specific sizes. - # At runtime, token counts not in this list fall back to eager execution. - # Example: "1,2,4,8,16,32,64,128,256,512" - "FD_BLOCK_WISE_CUDA_GRAPH_SIZES": lambda: os.getenv("FD_BLOCK_WISE_CUDA_GRAPH_SIZES", "128,256,512,1024,2048"), # Suspend rollouting routing replay "FD_SUSPEND_ROUTING_REPLAY": lambda: bool(int(os.getenv("FD_SUSPEND_ROUTING_REPLAY", "0"))), # train-infer consistency, used in RL diff --git a/fastdeploy/model_executor/graph_optimization/cuda_graph_op.py b/fastdeploy/model_executor/graph_optimization/cuda_graph_op.py deleted file mode 100644 index 3daecc87f69..00000000000 --- a/fastdeploy/model_executor/graph_optimization/cuda_graph_op.py +++ /dev/null @@ -1,320 +0,0 @@ -""" -# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License" -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" - -import functools -import inspect -from typing import Callable, Optional, Sequence - -import paddle -from paddleformers.utils.log import logger as _LOGGER - -import fastdeploy - -# ---- Module-level state for pre-captured block-wise CUDA graphs ---- - -# When True, the wrapper is in the capture phase (during dummy_run) and -# will capture new graphs. When False, uncached keys fall back to eager. -_BLOCK_WISE_CAPTURING: bool = False - -# Registry of all shared-mode graph caches, for bulk clearing. -_ALL_SHARED_CACHES: list = [] - -# Global counter / registry of all captured block-wise graphs (for logging). -# Each entry: (qualname, key, shared_mode) -_CAPTURED_GRAPH_LOG: list = [] - - -def get_captured_graph_log(): - """Return the list of all captured (qualname, key, shared) triples.""" - return list(_CAPTURED_GRAPH_LOG) - - -def dump_captured_graph_summary(): - """Print a summary of all captured block-wise CUDA graphs.""" - from collections import Counter - - if not fastdeploy.envs.FD_BLOCK_WISE_DEBUG: - return - if not _CAPTURED_GRAPH_LOG: - _LOGGER.info("[block_wise_cuda_graph] no graph captured") - return - counter = Counter(q for q, _, _ in _CAPTURED_GRAPH_LOG) - _LOGGER.info( - f"[block_wise_cuda_graph] total captured graphs={len(_CAPTURED_GRAPH_LOG)} " - f"across {len(counter)} distinct methods:" - ) - for qname, cnt in sorted(counter.items(), key=lambda x: -x[1]): - _LOGGER.info(f" - {qname} : {cnt} graph(s)") - - -def set_block_wise_capturing(capturing: bool): - """Toggle the capture phase flag. Only capture graphs when this is True.""" - global _BLOCK_WISE_CAPTURING - _BLOCK_WISE_CAPTURING = capturing - - -def clear_all_block_wise_graphs(): - """Clear all shared block-wise graph caches (e.g. for RL weight updates).""" - for graphs, cinputs, coutputs in _ALL_SHARED_CACHES: - graphs.clear() - cinputs.clear() - coutputs.clear() - - -def block_wise_cuda_graph_wrap( - inputs: Sequence[str], - self_attrs: Sequence[str] = (), - key_fn: Optional[Callable[..., tuple]] = None, -): - """ - Method decorator that wraps a forward method with CUDA Graph capture/replay. - - On the first call for a given cache key (derived from tensor shapes/dtypes), - the decorated method is captured into a CUDA Graph. Subsequent calls with the - same key will replay the graph after updating input data pointers. - - When ``_BLOCK_WISE_CAPTURING`` is managed via ``set_block_wise_capturing``, - new graphs are only captured during the capture phase (dummy_run). At runtime, - uncached keys fall back to eager execution, avoiding expensive on-the-fly captures. - - When ``self_attrs`` is provided, the named tensor attributes of ``self`` - (e.g. ``weight``) are also tracked for pointer replacement, and the graph - cache is **shared across all instances** (closure-level). This allows layers - with identical computation but different weights to share a single captured - graph, dramatically reducing the total number of graphs from O(num_layers) - to O(num_unique_shapes). - - When ``self_attrs`` is empty (default), graphs are cached per instance. - - Output tensors from the capture phase are reused across replays — the graph - always writes to the same output memory. This avoids per-replay allocation - overhead. Callers must consume the output before the next replay of the same - graph (which is naturally satisfied in sequential layer-by-layer forward). - - Args: - inputs: Names of parameters that are input tensors to be tracked for - CUDA Graph pointer replacement. These must be parameter names of the - decorated method. Only non-None tensor arguments are tracked. - self_attrs: Attribute names on ``self`` that are tensor parameters to be - replaced via pointer replacement (e.g. ``["weight"]``). When non-empty, - enables cross-instance graph sharing. - key_fn: Optional callable to generate the cache key from method arguments. - Signature: key_fn(arg0, arg1, ...) with args in declaration order - (excluding self). Defaults to a key based on tensor shapes/dtypes. - - Example: - class MyNorm(nn.Layer): - @block_wise_cuda_graph_wrap( - inputs=["x", "residual"], - self_attrs=["weight"], # all layers share one graph - ) - def forward(self, x, residual=None): - return rms_norm(x, self.weight), residual - """ - - def decorator(method: Callable) -> Callable: - sig = inspect.signature(method) - params = list(sig.parameters.keys()) # ["self", "x", "residual_input", ...] - _qualname = method.__qualname__ - - for name in inputs: - if name not in params or name == "self": - raise ValueError( - f"cuda_graph_wrap: input '{name}' is not a parameter of " - f"{method.__qualname__}. Available: {[p for p in params if p != 'self']}" - ) - - # ---- Pre-compute at decoration time (runs once) ---- - - _EMPTY = inspect.Parameter.empty - _Tensor = paddle.Tensor - - # For each non-self param: (name, args_index, default_value) - # args_index is position in *args (0-based, since self is consumed by Python) - _param_info = tuple((p, i - 1, sig.parameters[p].default) for i, p in enumerate(params) if p != "self") - - # For each declared input tensor: (name, args_index) - _input_info = tuple((name, params.index(name) - 1) for name in inputs) - - _self_attr_names = tuple(self_attrs) - _shared = len(_self_attr_names) > 0 - - _use_custom_key = key_fn is not None - - # --- Cache storage --- - # When self_attrs is provided: closure-level (shared across all instances) - # When not: per-instance (stored in self.__dict__) - if _shared: - _shared_graphs = {} - _shared_cinputs = {} - _shared_coutputs = {} # stores actual result tensors (reused across replays) - _ALL_SHARED_CACHES.append((_shared_graphs, _shared_cinputs, _shared_coutputs)) - - # Per-instance attribute key names - _g = f"_cg_{method.__name__}_g" - _ci = f"_cg_{method.__name__}_ci" - _co = f"_cg_{method.__name__}_co" - - @functools.wraps(method) - def wrapper(self, *args, **kwargs): - if not fastdeploy.envs.FD_USE_BLOCK_WISE_CUDA_GRAPH: - return method(self, *args, **kwargs) - - nargs = len(args) - - # Skip CUDA graph if any input tensor has a 0 in its shape - for a in args: - if isinstance(a, _Tensor) and 0 in a.shape: - return method(self, *args, **kwargs) - for v in kwargs.values(): - if isinstance(v, _Tensor) and 0 in v.shape: - return method(self, *args, **kwargs) - - # === Key generation: inline, no sig.bind === - if _use_custom_key: - # Resolve all args for custom key_fn - resolved = [] - for pname, aidx, default in _param_info: - if pname in kwargs: - resolved.append(kwargs[pname]) - elif aidx < nargs: - resolved.append(args[aidx]) - elif default is not _EMPTY: - resolved.append(default) - else: - resolved.append(None) - key = key_fn(*resolved) - else: - # Default: fast inline key from shapes/dtypes - _kp = [] - for pname, aidx, default in _param_info: - if pname in kwargs: - v = kwargs[pname] - elif aidx < nargs: - v = args[aidx] - else: - v = default - if isinstance(v, _Tensor): - _kp.append((tuple(v.shape), v.dtype)) - elif v is None: - _kp.append(None) - elif callable(v): - _kp.append(True) - # Include self_attrs shapes/dtypes in key - for attr_name in _self_attr_names: - attr = getattr(self, attr_name, None) - if attr is not None and isinstance(attr, _Tensor): - _kp.append((attr_name, tuple(attr.shape), attr.dtype)) - else: - _kp.append((attr_name, None)) - key = tuple(_kp) - - # === Get cache (shared or per-instance) === - if _shared: - graphs = _shared_graphs - cinputs = _shared_cinputs - coutputs = _shared_coutputs - else: - _d = self.__dict__ - try: - graphs = _d[_g] - cinputs = _d[_ci] - coutputs = _d[_co] - except KeyError: - graphs = {} - cinputs = {} - coutputs = {} - _d[_g] = graphs - _d[_ci] = cinputs - _d[_co] = coutputs - - if key not in graphs: - # === First encounter: only capture during capture phase === - if not _BLOCK_WISE_CAPTURING: - # Not in capture phase -- fall back to eager - return method(self, *args, **kwargs) - - # === Capture === - graph = paddle.device.cuda.graphs.CUDAGraph(enable_replace=True) - graphs[key] = graph - ci = {} - for name, aidx in _input_info: - v = kwargs[name] if name in kwargs else (args[aidx] if aidx < nargs else None) - if v is not None and isinstance(v, _Tensor): - ci[name] = v.data_ptr() - - # Record self_attrs pointers for cross-instance replacement - for attr_name in _self_attr_names: - attr = getattr(self, attr_name, None) - if attr is not None and isinstance(attr, _Tensor): - ci[f"__attr_{attr_name}"] = attr.data_ptr() - - cinputs[key] = ci - - graph.capture_begin() - result = method(self, *args, **kwargs) - graph.capture_end() - - # --- Log which op just entered the CUDA graph --- - _CAPTURED_GRAPH_LOG.append((_qualname, key, _shared)) - if fastdeploy.envs.FD_BLOCK_WISE_DEBUG: - _LOGGER.info( - f"[block_wise_cuda_graph] captured #{len(_CAPTURED_GRAPH_LOG)} " - f"op={_qualname} shared={_shared} key={key}" - ) - - graph.replay() - - # Store the actual result for reuse. The graph always writes to - # the same output memory, so we return the same tensors on replay. - coutputs[key] = result - return result - else: - # === Replay path (HOT PATH) === - old_ptrs = [] - new_ptrs = [] - ci = cinputs[key] - - for name, aidx in _input_info: - v = kwargs[name] if name in kwargs else (args[aidx] if aidx < nargs else None) - if v is not None and name in ci: - old_ptrs.append(ci[name]) - new_ptr = v.data_ptr() - new_ptrs.append(new_ptr) - ci[name] = new_ptr - - # Replace self_attrs pointers (e.g. weight) - for attr_name in _self_attr_names: - attr_key = f"__attr_{attr_name}" - if attr_key in ci: - attr = getattr(self, attr_name, None) - if attr is not None: - old_ptrs.append(ci[attr_key]) - new_ptr = attr.data_ptr() - new_ptrs.append(new_ptr) - ci[attr_key] = new_ptr - - if old_ptrs: - graphs[key].replace_input_ptrs(old_ptrs, new_ptrs) - graphs[key].replay() - - # Reuse the output tensors from capture — graph wrote fresh - # data to the same memory, no allocation needed. - return coutputs[key] - - return wrapper - - return decorator diff --git a/fastdeploy/model_executor/graph_optimization/cudagraph_piecewise_backend.py b/fastdeploy/model_executor/graph_optimization/cudagraph_piecewise_backend.py index d2d5a6993a6..1b5e73e4d8f 100644 --- a/fastdeploy/model_executor/graph_optimization/cudagraph_piecewise_backend.py +++ b/fastdeploy/model_executor/graph_optimization/cudagraph_piecewise_backend.py @@ -193,19 +193,14 @@ def run_static_model(self, entry: ConcreteSizeEntry, is_decode: bool = False, ** def __call__(self, **kwargs) -> List[paddle.Tensor] | paddle.Tensor: # Get real shape (total num tokens) - if ( - self.speculative_decoding - and self.real_bsz_to_captured_size - and all(self.real_bsz_to_captured_size.values()) - ): - seq_lens_this_time: paddle.Tensor = kwargs["forward_meta"].seq_lens_this_time - real_bsz = kwargs["forward_meta"].real_bsz - num_running_requests = real_bsz if real_bsz > 0 else int((seq_lens_this_time.flatten() > 0).sum().item()) - num_running_requests = max(1, num_running_requests) - real_shape = self.real_bsz_to_captured_size[num_running_requests] - else: - ids_remove_padding: paddle.Tensor = kwargs["forward_meta"].ids_remove_padding - real_shape = ids_remove_padding.shape[0] + # For both MTP speculative decoding and regular decode, use ids_remove_padding.shape[0] + # directly as the real_shape key into real_shape_to_captured_size. + # In MTP, cudagraph_capture_sizes are already scaled by (num_speculative_tokens+1), + # so ids_remove_padding.shape[0] == capture_size and the lookup is correct. + # This avoids using real_bsz (a concrete Python int) which causes SOT to specialize + # on each distinct batch size, generating multiple code objects and breaking warmup_impl. + ids_remove_padding: paddle.Tensor = kwargs["forward_meta"].ids_remove_padding + real_shape = ids_remove_padding.shape[0] exist_prefill = kwargs["forward_meta"].exist_prefill # Static split graph mode: use Static + CUDAGraph for prefill/mixed phase static_cudagraph_for_prefill = exist_prefill and not self.full_cuda_graph and self.dy2st diff --git a/fastdeploy/model_executor/graph_optimization/utils.py b/fastdeploy/model_executor/graph_optimization/utils.py index 2b5241e5aed..8d38df5324d 100644 --- a/fastdeploy/model_executor/graph_optimization/utils.py +++ b/fastdeploy/model_executor/graph_optimization/utils.py @@ -140,3 +140,9 @@ def get_state(): sot_warmup_guard, in_sot_warmup_mode = create_guard(False) profile_run_guard, in_profile_run_mode = create_guard(False) + +# Guard for the prefill SOT warmup + piecewise-CUDAGraph capture phase. +# When active, keep the SOT-compiled graph intact during prefill capture +# by avoiding extra graph break points or graph fragmentation. +# The decode-only CUDAGraph capture runs OUTSIDE this guard. +prefill_cudagraph_guard, in_prefill_cudagraph_mode = create_guard(False) diff --git a/fastdeploy/model_executor/layers/attention/append_attn_backend.py b/fastdeploy/model_executor/layers/attention/append_attn_backend.py index 594dbfe3763..f3941946e7e 100644 --- a/fastdeploy/model_executor/layers/attention/append_attn_backend.py +++ b/fastdeploy/model_executor/layers/attention/append_attn_backend.py @@ -331,17 +331,6 @@ def forward_mixed( q_norm_weight = getattr(layer, "q_norm_weight", None) if norm_after_rope_in_kernel else None k_norm_weight = getattr(layer, "k_norm_weight", None) if norm_after_rope_in_kernel else None - if self.rope_3d: - assert len(forward_meta.rotary_embs.shape) == 6 - else: - assert len(forward_meta.rotary_embs.shape) == 5 - if layer.use_neox_rotary_style: - assert forward_meta.rotary_embs.shape[0:4] == [2, 1, self.max_seq_len, 1] - # 128 is qwen3 - # 32 is glm - # 64 is gpt-oss - assert forward_meta.rotary_embs.shape[4] in [128, 32, 64] - if self.pd_disaggregation_mode == "per_query": metadata.kv_signal_data_list[layer.layer_id] = init_signal_layerwise( metadata.kv_signal_metadata, diff --git a/fastdeploy/model_executor/layers/attention/flash_attn_backend.py b/fastdeploy/model_executor/layers/attention/flash_attn_backend.py index fd55046051a..7c8a24c6005 100644 --- a/fastdeploy/model_executor/layers/attention/flash_attn_backend.py +++ b/fastdeploy/model_executor/layers/attention/flash_attn_backend.py @@ -56,6 +56,7 @@ ) from fastdeploy.model_executor.layers.attention.utils import init_rank_and_device_id from fastdeploy.model_executor.utils import get_sm_version +from fastdeploy.utils import register_custom_python_op if TYPE_CHECKING: from fastdeploy.model_executor.forward_meta import ForwardMeta @@ -229,7 +230,7 @@ class FlashAttentionMetadata(AttentionMetadata): FlashAttentionMetadata """ - cu_seqlens_k: paddle.Tensor = None + cu_seqlens_k: Optional[paddle.Tensor] = None pre_cache_batch_ids = None pre_cache_tile_ids_per_batch = None @@ -243,9 +244,469 @@ class FlashAttentionMetadata(AttentionMetadata): _fuse_kernel_compute_dtype: str = "bf16" _dtype: paddle.dtype = paddle.bfloat16 - max_len_tensor_cpu_decoder: paddle.Tensor = None + max_len_tensor_cpu_decoder: Optional[paddle.Tensor] = None - attn_mask_q: paddle.Tensor = None + attn_mask_q: Optional[paddle.Tensor] = None + + +class FlashAttnLayerCtx: + """ + Container for optional/variable Tensor args and scalar constants for + python_op_flash_attn_forward. Passed as a single const (non-Tensor) + parameter so that register_custom_python_op never sees it as a mutable + pir.Value regardless of whether the fields are None or real Tensors. + This pattern follows python_op_fused_moe_kernel_paddle's quant_config arg. + The object is created once in FlashAttentionBackend.__init__ and its + attributes are updated in-place each call, so SOT trace never inlines + __init__ and treats it as a stable ObjectVariable constant. + """ + + +def _python_op_flash_attn_forward_infer_meta( + qkv, + cache_k, + cache_v, + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + cu_seqlens_k, + block_tables, + encoder_batch_ids, + encoder_tile_ids_per_batch, + encoder_num_blocks_x_cpu, + kv_batch_ids, + kv_tile_ids_per_batch, + kv_num_blocks_x_cpu, + decoder_batch_ids, + decoder_tile_ids_per_batch, + decoder_num_blocks_cpu, + decoder_num_blocks_device, + decoder_chunk_size_device, + max_len_tensor_cpu, + rotary_embs, + layer_ctx, +): + token_num = qkv.shape[0] + return paddle.static.MetaTensor(shape=[token_num, layer_ctx.num_heads * layer_ctx.head_dim], dtype=qkv.dtype) + + +@register_custom_python_op( + name="python_op_flash_attn_forward", + infer_meta=_python_op_flash_attn_forward_infer_meta, + input_names=[ + "qkv", + "cache_k", + "cache_v", + "seq_lens_encoder", + "seq_lens_decoder", + "seq_lens_this_time", + "batch_id_per_token", + "cu_seqlens_q", + "cu_seqlens_k", + "block_tables", + "encoder_batch_ids", + "encoder_tile_ids_per_batch", + "encoder_num_blocks_x_cpu", + "kv_batch_ids", + "kv_tile_ids_per_batch", + "kv_num_blocks_x_cpu", + "decoder_batch_ids", + "decoder_tile_ids_per_batch", + "decoder_num_blocks_cpu", + "decoder_num_blocks_device", + "decoder_chunk_size_device", + "max_len_tensor_cpu", + "rotary_embs", + ], + output_names=["out"], + inplace_map={}, +) +def python_op_flash_attn_forward( + qkv, + cache_k, + cache_v, + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + cu_seqlens_k, + block_tables, + encoder_batch_ids, + encoder_tile_ids_per_batch, + encoder_num_blocks_x_cpu, + kv_batch_ids, + kv_tile_ids_per_batch, + kv_num_blocks_x_cpu, + decoder_batch_ids, + decoder_tile_ids_per_batch, + decoder_num_blocks_cpu, + decoder_num_blocks_device, + decoder_chunk_size_device, + max_len_tensor_cpu, + rotary_embs, + layer_ctx, +): + """ + Wraps FlashAttentionBackend forward_mixed as a single py_op so SOT treats + it as a black box (no BreakGraphError from int(TensorVariable)). The op is + placed in FLAGS_cuda_graph_blacklist and runs eagerly outside the CUDA + graph every step. + + get_block_shape_and_split_kv_block is called INSIDE this py-op (guarded by + layer_id == 0 so it runs once per step). It must run here rather than as a + separately-blacklisted custom op because it conditionally skips the DtoH + copy that refreshes max_len_tensor_cpu when IsCUDAGraphCapturing() is true + (see get_block_shape_and_split_kv_block.cu). Only by running fully eagerly + inside this py-op is IsCUDAGraphCapturing() guaranteed false at replay, so + max_len_tensor_cpu / decoder_num_blocks_cpu are always fresh; otherwise + downstream pre_cache_len_concat would size buffers from a stale (warmup) + max_dec_len and write out of bounds (CUDA illegal memory access). + """ + cache_k_scales = layer_ctx.cache_k_scales + cache_v_scales = layer_ctx.cache_v_scales + cache_k_out_scale = layer_ctx.cache_k_out_scale + cache_v_out_scale = layer_ctx.cache_v_out_scale + cache_k_zp = layer_ctx.cache_k_zp + cache_v_zp = layer_ctx.cache_v_zp + attn_mask = layer_ctx.attn_mask + attn_mask_offsets = layer_ctx.attn_mask_offsets + qkv_bias = layer_ctx.qkv_bias + qkv_scale = layer_ctx.qkv_scale + linear_shift = layer_ctx.linear_shift + linear_smooth = layer_ctx.linear_smooth + q_norm_weight = layer_ctx.q_norm_weight + k_norm_weight = layer_ctx.k_norm_weight + kv_signal_data = layer_ctx.kv_signal_data + sinks = layer_ctx.sinks + decode_block_indices = layer_ctx.decode_block_indices + decode_num_blocks = layer_ctx.decode_num_blocks + decode_chunk_size = layer_ctx.decode_chunk_size + decode_tmp_workspace = layer_ctx.decode_tmp_workspace + decode_tmp_m = layer_ctx.decode_tmp_m + decode_tmp_d = layer_ctx.decode_tmp_d + num_heads = layer_ctx.num_heads + kv_num_heads = layer_ctx.kv_num_heads + head_dim = layer_ctx.head_dim + attn_outputsize_tp = layer_ctx.attn_outputsize_tp + max_seq_len = layer_ctx.max_seq_len + encoder_block_shape_q = layer_ctx.encoder_block_shape_q + decoder_block_shape_q = layer_ctx.decoder_block_shape_q + group_size = layer_ctx.group_size + block_size = layer_ctx.block_size + max_partition_size = layer_ctx.max_partition_size + max_tokens_per_batch = layer_ctx.max_tokens_per_batch + speculate_max_draft_token_num = layer_ctx.speculate_max_draft_token_num + causal = layer_ctx.causal + use_speculate = layer_ctx.use_speculate + rope_3d = layer_ctx.rope_3d + fuse_kernel_compute_dtype = layer_ctx.fuse_kernel_compute_dtype + use_decode_unified_attention = layer_ctx.use_decode_unified_attention + rms_norm_eps = layer_ctx.rms_norm_eps + cache_quant_type_str = layer_ctx.cache_quant_type_str + use_neox_rotary_style = layer_ctx.use_neox_rotary_style + quant_max_bound = layer_ctx.quant_max_bound + quant_min_bound = layer_ctx.quant_min_bound + out_scale = layer_ctx.out_scale + + # get_block_shape_and_split_kv_block writes the per-step attention metadata + # (decoder/encoder/kv split buffers + max_len_tensor_cpu) in-place. Run it + # once per step (layer 0) fully eagerly here so the guarded DtoH copies that + # refresh max_len_tensor_cpu / decoder_num_blocks_cpu always execute + # (IsCUDAGraphCapturing() is false inside this eager py-op). Layers 1..N read + # the buffers written by layer 0. + if layer_ctx.layer_id == 0: + get_block_shape_and_split_kv_block( + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + decoder_batch_ids, + decoder_tile_ids_per_batch, + decoder_num_blocks_cpu, + decoder_num_blocks_device, + decoder_chunk_size_device, + max_len_tensor_cpu, + encoder_batch_ids, + encoder_tile_ids_per_batch, + encoder_num_blocks_x_cpu, + kv_batch_ids, + kv_tile_ids_per_batch, + kv_num_blocks_x_cpu, + encoder_block_shape_q, + decoder_block_shape_q, + group_size, + block_size, + ) + + # NOTE: This metadata (pre_cache_len_concat / get_attn_mask_q / config_for_attention) + # depends only on tensor inputs that are identical across all layers within a step. + # It is computed unconditionally here (rather than cached at layer_id==0 on layer_ctx) + # because under piecewise cudagraph the py-op body runs at graph-execution time while + # layer_ctx (a shared const object) only ever holds its last-traced values, which would + # make any cross-layer caching stale. Per-layer recompute is cheap and correct. + # get_block_shape_and_split_kv_block runs eagerly inside this py-op at layer 0 + # (IsCUDAGraphCapturing() is false here), so its blocking DtoH copy DOES refresh + # max_len_tensor_cpu with the real per-step values (see get_block_shape_...cu:296). + # Read the host scalars straight from it, exactly like the eager (non-cudagraph) path. + # + # Do NOT recompute these from the live seq_lens tensors: index [3] must be + # max(seq_len_decoder + seq_len_this_time) over slots with seq_len_this_time > 0 + # (GetMaxLenKernel skips seq_len_this_time <= 0). A naive + # (seq_lens_decoder + seq_lens_this_time).max() over ALL slots picks up stale + # seq_lens_decoder in inactive/padded slots -> too-large max_seqlen_k -> the FA3 + # varlen kernel schedules K tiles past the end of k/v -> illegal memory access + # (flash_fwd_launch_template.h:160). + max_len_cpu = max_len_tensor_cpu + use_fa_do_prefill = int(max_len_cpu[1]) > 0 + if use_fa_do_prefill: + max_len_tensor_cpu_decoder = paddle.clone(max_len_tensor_cpu) + max_len_tensor_cpu_decoder[1] = 0 + ( + cu_seqlens_k, + pre_cache_batch_ids, + pre_cache_tile_ids_per_batch, + pre_cache_num_blocks_cpu, + kv_token_num_cpu, + ) = pre_cache_len_concat( + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + max_len_cpu[2], + block_size, + ) + if FLASH_ATTN_VERSION == 4 or attn_mask_offsets is not None: + attn_mask_q = get_attn_mask_q( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + attn_mask_kv=attn_mask_offsets, + kv_token_num=int(kv_token_num_cpu[0]), + ) + else: + attn_mask_q = None + else: + max_len_tensor_cpu_decoder = max_len_tensor_cpu + pre_cache_batch_ids = None + pre_cache_tile_ids_per_batch = None + pre_cache_num_blocks_cpu = None + kv_token_num_cpu = None + attn_mask_q = None + if use_decode_unified_attention: + config_for_attention( + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + decode_block_indices, + decode_num_blocks, + decode_chunk_size, + max_len_tensor_cpu, + cache_quant_type_str, + group_size, + kv_num_heads, + max_tokens_per_batch, + ) + + token_num = qkv.shape[0] + batch_id_per_token = batch_id_per_token.flatten() + real_len = batch_id_per_token.shape[0] + if real_len < token_num: + batch_id_per_token = paddle.nn.functional.pad(batch_id_per_token, [0, token_num - real_len], value=-1) + elif real_len > token_num: + batch_id_per_token = batch_id_per_token[:token_num] + + if use_fa_do_prefill: + q, k, v, _ = gqa_rope_write_cache( + qkv, + cache_k, + cache_v, + cu_seqlens_q, + cu_seqlens_k, + rotary_embs, + seq_lens_this_time, + seq_lens_encoder, + seq_lens_decoder, + batch_id_per_token, + block_tables, + kv_batch_ids, + kv_tile_ids_per_batch, + kv_num_blocks_x_cpu, + pre_cache_batch_ids, + pre_cache_tile_ids_per_batch, + pre_cache_num_blocks_cpu, + q_norm_weight, + k_norm_weight, + cache_k_scales, + cache_v_scales, + cache_k_out_scale, + cache_v_out_scale, + cache_k_zp, + cache_v_zp, + kv_signal_data, + int(kv_token_num_cpu[0]), + max_seq_len, + rms_norm_eps, + use_neox_rotary_style, + cache_quant_type_str, + rope_3d, + ) + res_encoder = flash_attn_func( + q, + k, + v, + cu_seqlens_q[: cu_seqlens_k.shape[0]], + cu_seqlens_k, + max_seqlen_q=max_len_cpu[0], + max_seqlen_k=max_len_cpu[3], + attn_mask_q=attn_mask_q, + causal=causal, + num_heads=num_heads, + kv_num_heads=kv_num_heads, + head_dim=head_dim, + )[0].reshape([-1, attn_outputsize_tp]) + + if use_decode_unified_attention: + qkv_out = decoder_write_cache_with_rope( + qkv, + cache_k, + cache_v, + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + block_tables, + max_len_tensor_cpu, + rotary_embs, + qkv_bias, + cache_k_scales, + cache_v_scales, + cache_k_out_scale, + cache_v_out_scale, + cache_k_zp, + cache_v_zp, + kv_signal_data, + q_norm_weight, + k_norm_weight, + rms_norm_eps, + cache_quant_type_str, + use_neox_rotary_style, + rope_3d, + max_seq_len, + quant_max_bound, + quant_min_bound, + use_speculate, + ) + res_decoder = paddle.empty([qkv.shape[0], num_heads * head_dim], dtype=qkv.dtype) + if use_fa_do_prefill: + res_decoder[: res_encoder.shape[0]] = res_encoder + decode_unified_attention( + qkv_out, + cache_k, + cache_v, + decode_tmp_workspace, + decode_tmp_m, + decode_tmp_d, + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + block_tables, + decode_block_indices, + decode_num_blocks, + decode_chunk_size, + max_len_tensor_cpu, + attn_mask, + cache_k_scales, + cache_v_scales, + cache_k_out_scale, + cache_v_out_scale, + cache_k_zp, + cache_v_zp, + attn_mask_offsets, + sinks, + res_decoder, + cache_quant_type_str, + max_seq_len, + quant_max_bound, + quant_min_bound, + speculate_max_draft_token_num + 1, + causal, + ) + return res_decoder + else: + res_decoder = append_attention( + qkv, + cache_k, + cache_v, + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + batch_id_per_token, + cu_seqlens_q, + block_tables, + encoder_batch_ids, + encoder_tile_ids_per_batch, + encoder_num_blocks_x_cpu, + kv_batch_ids, + kv_tile_ids_per_batch, + kv_num_blocks_x_cpu, + decoder_batch_ids, + decoder_tile_ids_per_batch, + decoder_num_blocks_cpu, + max_len_tensor_cpu_decoder, + rotary_embs, + attn_mask, + qkv_bias, + qkv_scale, + cache_k_scales, + cache_v_scales, + cache_k_out_scale, + cache_v_out_scale, + cache_k_zp, + cache_v_zp, + linear_shift, + linear_smooth, + attn_mask_offsets, + kv_signal_data, + q_norm_weight, + k_norm_weight, + sinks, + rms_norm_eps, + fuse_kernel_compute_dtype, + cache_quant_type_str, + use_neox_rotary_style, + rope_3d, + max_seq_len, + quant_max_bound, + quant_min_bound, + out_scale, + encoder_block_shape_q, + decoder_block_shape_q, + max_partition_size, + max_seq_len, + speculate_max_draft_token_num + 1, + causal, + use_speculate, + ) + if use_fa_do_prefill: + merge_prefill_decode_output( + res_encoder, + res_decoder, + seq_lens_encoder, + seq_lens_decoder, + seq_lens_this_time, + cu_seqlens_q, + num_heads, + head_dim, + speculate_max_draft_token_num + 1, + ) + attn_out = paddle.empty([qkv.shape[0], num_heads * head_dim], dtype=qkv.dtype) + attn_out[: res_encoder.shape[0]] = res_encoder + return attn_out + else: + return res_decoder class FlashAttentionBackend(AttentionBackend): @@ -305,6 +766,43 @@ def __init__( if FLASH_ATTN_VERSION is None: init_flash_attn_version() + # In static split-graph (piecewise cudagraph) mode, the dynamic attention + # region cannot be captured. python_op_flash_attn_forward wraps the whole + # dynamic logic (including get_block_shape_and_split_kv_block) into a single + # py-op; blacklist it so it runs eagerly outside the CUDA graph. Running + # get_block_shape inside this eager py-op (rather than as a separate + # blacklisted custom op) guarantees IsCUDAGraphCapturing() is false when its + # guarded DtoH copies refresh max_len_tensor_cpu / decoder_num_blocks_cpu. + if not fd_config.graph_opt_config.full_cuda_graph: + flag = "FLAGS_cuda_graph_blacklist" + paddle.set_flags( + { + flag: ",".join( + list( + set( + paddle.get_flags(flag)[flag].split(",") + + [ + "py_op.python_op_flash_attn_forward_", + ] + ) + ) + ) + } + ) + + # Per-layer, id-stable context objects carrying that layer's non-Tensor + # scalars and (layer-static) weight tensors into python_op_flash_attn_forward + # as one const param. IMPORTANT: there must be a DISTINCT object per layer. + # + # Under piecewise cudagraph the traced forward runs ONCE; python_op_flash_attn_forward's + # body then runs eagerly each step reading the const object by reference. A single + # shared object would freeze to the LAST layer's values at trace time, so every + # layer's node would read layer N-1's weights. Keying by layer_id gives each node its + # own frozen-but-correct object (layer weights are static across steps). The map is + # populated lazily at trace time and reused across re-traces (stable ids -> stable + # register_op specializations). + self.layer_ctxs: dict = {} + def get_attention_meta(self): """get_attention_meta""" return self.attention_metadata @@ -394,266 +892,82 @@ def forward_mixed( cache_k_scales = getattr(layer, "cache_k_scale", None) cache_v_scales = getattr(layer, "cache_v_scale", None) - if layer.layer_id == 0: - get_block_shape_and_split_kv_block( - forward_meta.seq_lens_encoder, - forward_meta.seq_lens_decoder, - forward_meta.seq_lens_this_time, - forward_meta.decoder_batch_ids, - forward_meta.decoder_tile_ids_per_batch, - forward_meta.decoder_num_blocks_cpu, - forward_meta.decoder_num_blocks_device, - forward_meta.decoder_chunk_size_device, - forward_meta.max_len_tensor_cpu, - forward_meta.encoder_batch_ids, - forward_meta.encoder_tile_ids_per_batch, - forward_meta.encoder_num_blocks_x_cpu, - forward_meta.kv_batch_ids, - forward_meta.kv_tile_ids_per_batch, - forward_meta.kv_num_blocks_x_cpu, - self.encoder_block_shape_q, - self.decoder_block_shape_q, - self.group_size, - self.block_size, - ) - - if forward_meta.max_len_tensor_cpu[1].item() > 0: - - forward_meta.max_len_tensor_cpu_decoder = paddle.clone(forward_meta.max_len_tensor_cpu) - forward_meta.max_len_tensor_cpu_decoder[1] = 0 - - ( - forward_meta.cu_seqlens_k, - forward_meta.pre_cache_batch_ids, - forward_meta.pre_cache_tile_ids_per_batch, - forward_meta.pre_cache_num_blocks_cpu, - forward_meta.kv_token_num_cpu, - ) = pre_cache_len_concat( - forward_meta.seq_lens_encoder, - forward_meta.seq_lens_decoder, - forward_meta.seq_lens_this_time, - forward_meta.max_len_tensor_cpu[2], - self.block_size, - ) - if FLASH_ATTN_VERSION == 4 or forward_meta.attn_mask_offsets is not None: - forward_meta.attn_mask_q = get_attn_mask_q( - cu_seqlens_q=forward_meta.cu_seqlens_q, - cu_seqlens_k=forward_meta.cu_seqlens_k, - attn_mask_kv=forward_meta.attn_mask_offsets, - kv_token_num=forward_meta.kv_token_num_cpu[0].item(), - ) - else: - forward_meta.attn_mask_q = None - if envs.USE_DECODE_UNIFIED_ATTENTION: - config_for_attention( - forward_meta.seq_lens_encoder, - forward_meta.seq_lens_decoder, - forward_meta.seq_lens_this_time, - forward_meta.decode_block_indices, - forward_meta.decode_num_blocks, - forward_meta.decode_chunk_size, - forward_meta.max_len_tensor_cpu, - getattr(layer, "cache_quant_type_str", "none"), - self.group_size, - self.kv_num_heads, - self.max_tokens_per_batch, - ) - - use_fa_do_prefill = forward_meta.max_len_tensor_cpu[1].item() > 0 - - if use_fa_do_prefill: - q, k, v, _ = gqa_rope_write_cache( - qkv, - cache_k, - cache_v, - forward_meta.cu_seqlens_q, - forward_meta.cu_seqlens_k, - forward_meta.rotary_embs, - forward_meta.seq_lens_this_time, - forward_meta.seq_lens_encoder, - forward_meta.seq_lens_decoder, - forward_meta.batch_id_per_token, - forward_meta.block_tables, - forward_meta.kv_batch_ids, - forward_meta.kv_tile_ids_per_batch, - forward_meta.kv_num_blocks_x_cpu, - forward_meta.pre_cache_batch_ids, - forward_meta.pre_cache_tile_ids_per_batch, - forward_meta.pre_cache_num_blocks_cpu, - q_norm_weight, - k_norm_weight, - cache_k_scales, - cache_v_scales, - getattr(layer, "cache_k_out_scale", None), - getattr(layer, "cache_v_out_scale", None), - getattr(layer, "cache_k_zp", None), - getattr(layer, "cache_v_zp", None), - metadata.kv_signal_data_list[layer.layer_id], - forward_meta.kv_token_num_cpu[0].item(), - self.max_seq_len, - getattr(layer, "rms_norm_eps", 1e-6), - layer.use_neox_rotary_style, - getattr(layer, "cache_quant_type_str", "none"), - self.rope_3d, - ) - - res_encoder = flash_attn_func( - q, - k, - v, - forward_meta.cu_seqlens_q[: forward_meta.cu_seqlens_k.shape[0]], - forward_meta.cu_seqlens_k, - max_seqlen_q=forward_meta.max_len_tensor_cpu[0], - max_seqlen_k=forward_meta.max_len_tensor_cpu[3], - attn_mask_q=forward_meta.attn_mask_q, - causal=self.causal, - num_heads=self.num_heads, - kv_num_heads=self.kv_num_heads, - head_dim=self.head_dim, - )[0].reshape([-1, self.attn_outputsize_tp]) - - if envs.USE_DECODE_UNIFIED_ATTENTION: - qkv_out = decoder_write_cache_with_rope( - qkv, - cache_k, - cache_v, - forward_meta.seq_lens_encoder, - forward_meta.seq_lens_decoder, - forward_meta.seq_lens_this_time, - forward_meta.batch_id_per_token, - forward_meta.cu_seqlens_q, - forward_meta.block_tables, - forward_meta.max_len_tensor_cpu, - forward_meta.rotary_embs, - layer.qkv_bias, - cache_k_scales, - cache_v_scales, - getattr(layer, "cache_k_out_scale", None), - getattr(layer, "cache_v_out_scale", None), - getattr(layer, "cache_k_zp", None), - getattr(layer, "cache_v_zp", None), - metadata.kv_signal_data_list[layer.layer_id], - q_norm_weight, - k_norm_weight, - getattr(layer, "rms_norm_eps", 1e-6), - getattr(layer, "cache_quant_type_str", "none"), - layer.use_neox_rotary_style, - self.rope_3d, - self.max_seq_len, - getattr(layer, "quant_max_bound", 0.0), - getattr(layer, "quant_min_bound", 0.0), - self.speculative_method is not None, - ) - if use_fa_do_prefill: - res_decoder = res_encoder - else: - res_decoder = paddle.empty( - [qkv.shape[0], self.num_heads * self.head_dim], - dtype=qkv.dtype, - ) - decode_unified_attention( - qkv_out, - cache_k, - cache_v, - forward_meta.decode_tmp_workspace, - forward_meta.decode_tmp_m, - forward_meta.decode_tmp_d, - forward_meta.seq_lens_encoder, - forward_meta.seq_lens_decoder, - forward_meta.seq_lens_this_time, - forward_meta.batch_id_per_token, - forward_meta.cu_seqlens_q, - forward_meta.block_tables, - forward_meta.decode_block_indices, - forward_meta.decode_num_blocks, - forward_meta.decode_chunk_size, - forward_meta.max_len_tensor_cpu, - forward_meta.attn_mask, - cache_k_scales, - cache_v_scales, - getattr(layer, "cache_k_out_scale", None), - getattr(layer, "cache_v_out_scale", None), - getattr(layer, "cache_k_zp", None), - getattr(layer, "cache_v_zp", None), - forward_meta.attn_mask_offsets, - getattr(layer, "sinks", None), - res_decoder, - getattr(layer, "cache_quant_type_str", "none"), - self.max_seq_len, - getattr(layer, "quant_max_bound", 0.0), - getattr(layer, "quant_min_bound", 0.0), - self.speculate_max_draft_token_num + 1, - self.causal, - ) - return res_decoder - else: - res_decoder = append_attention( - qkv, - cache_k, - cache_v, - forward_meta.seq_lens_encoder, - forward_meta.seq_lens_decoder, - forward_meta.seq_lens_this_time, - forward_meta.batch_id_per_token, - forward_meta.cu_seqlens_q, - forward_meta.block_tables, - forward_meta.encoder_batch_ids, - forward_meta.encoder_tile_ids_per_batch, - forward_meta.encoder_num_blocks_x_cpu, - forward_meta.kv_batch_ids, - forward_meta.kv_tile_ids_per_batch, - forward_meta.kv_num_blocks_x_cpu, - forward_meta.decoder_batch_ids, - forward_meta.decoder_tile_ids_per_batch, - forward_meta.decoder_num_blocks_cpu, - forward_meta.max_len_tensor_cpu_decoder if use_fa_do_prefill else forward_meta.max_len_tensor_cpu, - forward_meta.rotary_embs, - forward_meta.attn_mask, - layer.qkv_bias, - layer.qkv_scale, - cache_k_scales, - cache_v_scales, - getattr(layer, "cache_k_out_scale", None), - getattr(layer, "cache_v_out_scale", None), - getattr(layer, "cache_k_zp", None), - getattr(layer, "cache_v_zp", None), - layer.linear_shift, - layer.linear_smooth, - forward_meta.attn_mask_offsets, - metadata.kv_signal_data_list[layer.layer_id], - q_norm_weight, - k_norm_weight, - getattr(layer, "sinks", None), - getattr(layer, "rms_norm_eps", 1e-6), - metadata._fuse_kernel_compute_dtype, - getattr(layer, "cache_quant_type_str", "none"), - layer.use_neox_rotary_style, - self.rope_3d, - self.max_seq_len, - getattr(layer, "quant_max_bound", 0.0), - getattr(layer, "quant_min_bound", 0.0), - getattr(layer, "out_scale", -1.0), - self.encoder_block_shape_q, - self.decoder_block_shape_q, - self.max_partition_size, - self.max_seq_len, - self.speculate_max_draft_token_num + 1, - self.causal, - self.speculative_method is not None, - ) - - if use_fa_do_prefill: - merge_prefill_decode_output( - res_encoder, - res_decoder, - forward_meta.seq_lens_encoder, - forward_meta.seq_lens_decoder, - forward_meta.seq_lens_this_time, - forward_meta.cu_seqlens_q, - self.num_heads, - self.head_dim, - self.speculate_max_draft_token_num + 1, - ) - return res_encoder - else: - return res_decoder + ctx = self.layer_ctxs.get(layer.layer_id) + if ctx is None: + # First trace of this layer: create its own context and bake the + # layer-invariant scalars once. + ctx = FlashAttnLayerCtx() + ctx.num_heads = self.num_heads + ctx.kv_num_heads = self.kv_num_heads + ctx.head_dim = self.head_dim + ctx.attn_outputsize_tp = self.attn_outputsize_tp + ctx.max_seq_len = self.max_seq_len + ctx.encoder_block_shape_q = self.encoder_block_shape_q + ctx.decoder_block_shape_q = self.decoder_block_shape_q + ctx.group_size = self.group_size + ctx.block_size = self.block_size + ctx.max_partition_size = self.max_partition_size + ctx.max_tokens_per_batch = self.max_tokens_per_batch + ctx.speculate_max_draft_token_num = self.speculate_max_draft_token_num + ctx.causal = self.causal + ctx.use_speculate = self.use_speculate + ctx.rope_3d = self.rope_3d + self.layer_ctxs[layer.layer_id] = ctx + ctx.cache_k_scales = cache_k_scales + ctx.cache_v_scales = cache_v_scales + ctx.cache_k_out_scale = getattr(layer, "cache_k_out_scale", None) + ctx.cache_v_out_scale = getattr(layer, "cache_v_out_scale", None) + ctx.cache_k_zp = getattr(layer, "cache_k_zp", None) + ctx.cache_v_zp = getattr(layer, "cache_v_zp", None) + ctx.attn_mask = forward_meta.attn_mask + ctx.attn_mask_offsets = getattr(forward_meta, "attn_mask_offsets", None) + ctx.qkv_bias = layer.qkv_bias + ctx.qkv_scale = layer.qkv_scale + ctx.linear_shift = layer.linear_shift + ctx.linear_smooth = layer.linear_smooth + ctx.q_norm_weight = q_norm_weight + ctx.k_norm_weight = k_norm_weight + ctx.kv_signal_data = metadata.kv_signal_data_list[layer.layer_id] + ctx.sinks = getattr(layer, "sinks", None) + ctx.decode_block_indices = getattr(forward_meta, "decode_block_indices", None) + ctx.decode_num_blocks = getattr(forward_meta, "decode_num_blocks", None) + ctx.decode_chunk_size = getattr(forward_meta, "decode_chunk_size", None) + ctx.decode_tmp_workspace = getattr(forward_meta, "decode_tmp_workspace", None) + ctx.decode_tmp_m = getattr(forward_meta, "decode_tmp_m", None) + ctx.decode_tmp_d = getattr(forward_meta, "decode_tmp_d", None) + ctx.fuse_kernel_compute_dtype = metadata._fuse_kernel_compute_dtype + ctx.use_decode_unified_attention = envs.USE_DECODE_UNIFIED_ATTENTION + ctx.rms_norm_eps = getattr(layer, "rms_norm_eps", 1e-6) + ctx.cache_quant_type_str = cache_quant_type_str + ctx.use_neox_rotary_style = layer.use_neox_rotary_style + ctx.quant_max_bound = getattr(layer, "quant_max_bound", 0.0) + ctx.quant_min_bound = getattr(layer, "quant_min_bound", 0.0) + ctx.out_scale = getattr(layer, "out_scale", -1.0) + ctx.layer_id = layer.layer_id + + return python_op_flash_attn_forward( + qkv, + cache_k, + cache_v, + forward_meta.seq_lens_encoder, + forward_meta.seq_lens_decoder, + forward_meta.seq_lens_this_time, + forward_meta.batch_id_per_token, + forward_meta.cu_seqlens_q, + forward_meta.cu_seqlens_k, + forward_meta.block_tables, + forward_meta.encoder_batch_ids, + forward_meta.encoder_tile_ids_per_batch, + forward_meta.encoder_num_blocks_x_cpu, + forward_meta.kv_batch_ids, + forward_meta.kv_tile_ids_per_batch, + forward_meta.kv_num_blocks_x_cpu, + forward_meta.decoder_batch_ids, + forward_meta.decoder_tile_ids_per_batch, + forward_meta.decoder_num_blocks_cpu, + forward_meta.decoder_num_blocks_device, + forward_meta.decoder_chunk_size_device, + forward_meta.max_len_tensor_cpu, + forward_meta.rotary_embs, + layer_ctx=ctx, + ) diff --git a/fastdeploy/model_executor/layers/linear.py b/fastdeploy/model_executor/layers/linear.py index 193bd80cd90..b164031fc29 100644 --- a/fastdeploy/model_executor/layers/linear.py +++ b/fastdeploy/model_executor/layers/linear.py @@ -25,9 +25,6 @@ decode_alltoall_transpose, tensor_model_parallel_all_reduce, ) -from fastdeploy.model_executor.graph_optimization.cuda_graph_op import ( - block_wise_cuda_graph_wrap, -) from fastdeploy.model_executor.layers.quantization.quant_base import QuantMethodBase from fastdeploy.model_executor.utils import ( default_weight_loader, @@ -272,7 +269,6 @@ def load_state_dict(self, state_dict: dict): bias_tensor = paddle.to_tensor(get_tensor(state_dict.pop(self.bias_key))) self.bias.set_value(bias_tensor) - @block_wise_cuda_graph_wrap(inputs=["x"], self_attrs=["weight", "weight_scale_inv", "bias"]) def forward_cuda(self, x: paddle.Tensor) -> paddle.Tensor: """ Forward function for Linear. diff --git a/fastdeploy/model_executor/layers/moe/ep.py b/fastdeploy/model_executor/layers/moe/ep.py index 11c98463fcf..705424f43fe 100644 --- a/fastdeploy/model_executor/layers/moe/ep.py +++ b/fastdeploy/model_executor/layers/moe/ep.py @@ -27,6 +27,7 @@ import fastdeploy from fastdeploy import envs from fastdeploy.config import MoEPhase +from fastdeploy.model_executor.layers.moe.moe import get_moe_scores from fastdeploy.platforms import current_platform from fastdeploy.utils import singleton @@ -500,8 +501,6 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor): ) = layer.redundant_table_manger.get_ep_rank_to_expert_id_list_by_layer(layer.layer_idx) if layer.topk_method == "noaux_tc": - from .moe import get_moe_scores - score, topk_weights, topk_idx = get_moe_scores( gate_out, layer.n_group, @@ -530,8 +529,6 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor): ) else: if layer.topk_method == "noaux_tc": - from fastdeploy.model_executor.layers.moe.moe import get_moe_scores - use_fused = ( layer.fd_config.scheduler_config.enable_moe_scores_elementwise_fuse and current_platform.is_cuda() ) diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_triton_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_triton_backend.py index e3e9c5a0e95..cefa1a0c3d7 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_triton_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_triton_backend.py @@ -700,191 +700,33 @@ def apply( fc2_latent_proj: nn.Layer = None, ) -> paddle.Tensor: """ - Triton compute Fused MoE. + Triton compute Fused wfp8afp8 MoE via single blacklist custom op. """ - token_num = x.shape[0] - if token_num == 0: - return paddle.zeros([token_num, layer.hidden_size], dtype=x.dtype) gate_out = gate(x) - top_k = layer.top_k - num_local_experts = layer.num_local_experts - moe_intermediate_size = layer.moe_intermediate_size - hidden_size = layer.hidden_size - E, N1, _ = getattr(layer, self.added_weight_attrs[0]).shape - - if layer.topk_method == "noaux_tc": - use_fused = ( - layer.fd_config.scheduler_config.enable_moe_scores_elementwise_fuse and current_platform.is_cuda() - ) - if not use_fused: - gate_out = gate_out.cast("float32") - gate_out, topk_weights, topk_ids = get_moe_scores( - gate_out, - layer.n_group, - layer.topk_group, - layer.top_k, - layer.routed_scaling_factor, - layer.gate_correction_bias, - getattr(layer, "renormalize", True), - use_fused_cast=use_fused, - ) - else: - gate_out = gate_out.cast("float32") - topk_ids, topk_weights = fastdeploy.model_executor.ops.gpu.moe_topk_select( - gate_out, - layer.gate_correction_bias, - layer.top_k, - True, # apply_norm_weight - False, - ) + out = paddle.empty([x.shape[0], layer.hidden_size], dtype=x.dtype) - config = { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 256, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 32, - "num_warps": 8, - "num_stages": 4, - } - if token_num <= E: - config = { - "BLOCK_SIZE_M": 64, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 128, - "GROUP_SIZE_M": 1, - "num_warps": 4, - "num_stages": 4, - } - - sorted_token_ids, expert_ids, num_tokens_post_padded = tritonmoe_preprocess_func( - topk_ids, num_local_experts, config["BLOCK_SIZE_M"] - ) - max_possible_num_post_padded = sorted_token_ids.shape[0] - grid = ( - ceil_div(max_possible_num_post_padded, config["BLOCK_SIZE_M"]) - * ceil_div(moe_intermediate_size * 2, config["BLOCK_SIZE_N"]), - ) - - if topk_ids_hookfunc is not None: - topk_ids_hookfunc(topk_ids=topk_ids) - - up_gate_proj_out = paddle.empty( - [token_num * top_k, moe_intermediate_size * 2], - dtype=x.dtype, - ) - - from .triton_moe_kernels import fused_moe_kernel_paddle - - x_q, x_scale = scaled_fp8_quant(x, use_per_token_if_dynamic=True) - - fused_moe_kernel_paddle[grid]( - x_q, - layer.up_gate_proj_weight, - up_gate_proj_out, - x_scale, - layer.up_gate_proj_weight_scale, - None, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - max_possible_num_post_padded, - token_num * top_k, - N=moe_intermediate_size * 2, - K=hidden_size, - stride_am=x_q.strides[0], - stride_ak=x_q.strides[1], - stride_be=layer.up_gate_proj_weight.strides[0], - stride_bk=layer.up_gate_proj_weight.strides[2], - stride_bn=layer.up_gate_proj_weight.strides[1], - stride_cm=up_gate_proj_out.strides[0], - stride_cn=up_gate_proj_out.strides[1], - # - stride_asm=x_scale.strides[0], - stride_ask=x_scale.strides[1], - stride_bse=layer.up_gate_proj_weight_scale.strides[0], - stride_bsk=layer.up_gate_proj_weight_scale.strides[2], - stride_bsn=layer.up_gate_proj_weight_scale.strides[1], - group_n=-1, - group_k=-1, - # Meta-parameters - BLOCK_SIZE_M=config["BLOCK_SIZE_M"], - BLOCK_SIZE_N=config["BLOCK_SIZE_N"], - BLOCK_SIZE_K=config["BLOCK_SIZE_K"], - GROUP_SIZE_M=config["GROUP_SIZE_M"], - MUL_ROUTED_WEIGHT=False, - top_k=top_k, - compute_type_enum=1, - use_fp8_w8a8=True, - use_int8_w8a16=False, - per_channel_quant=True, - even_Ks=hidden_size % config["BLOCK_SIZE_K"] == 0, - num_warps=config.get("num_warps", 4), - num_stages=config.get("num_stages", 4), - ) - - down_proj_input = paddle.incubate.nn.functional.swiglu(up_gate_proj_out) - - down_proj_out = paddle.empty( - (token_num * top_k, hidden_size), - dtype=x.dtype, - ) - - grid = ( - ceil_div(max_possible_num_post_padded, config["BLOCK_SIZE_M"]) - * ceil_div(hidden_size, config["BLOCK_SIZE_N"]), - ) - - x_q, x_scale = scaled_fp8_quant(down_proj_input, use_per_token_if_dynamic=True) - - fused_moe_kernel_paddle[grid]( - x_q, - layer.down_proj_weight, - down_proj_out, - x_scale, - layer.down_proj_weight_scale, - topk_weights, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - max_possible_num_post_padded, - token_num * top_k, - N=hidden_size, - K=moe_intermediate_size, - stride_am=x_q.strides[0], - stride_ak=x_q.strides[1], - stride_be=layer.down_proj_weight.strides[0], - stride_bk=layer.down_proj_weight.strides[2], - stride_bn=layer.down_proj_weight.strides[1], - stride_cm=down_proj_out.strides[0], - stride_cn=down_proj_out.strides[1], - stride_asm=x_scale.strides[0], - stride_ask=x_scale.strides[1], - stride_bse=layer.down_proj_weight_scale.strides[0], - stride_bsk=layer.down_proj_weight_scale.strides[2], - stride_bsn=layer.down_proj_weight_scale.strides[1], - group_n=-1, - group_k=-1, - # Meta-parameters - BLOCK_SIZE_M=config["BLOCK_SIZE_M"], - BLOCK_SIZE_N=config["BLOCK_SIZE_N"], - BLOCK_SIZE_K=config["BLOCK_SIZE_K"], - GROUP_SIZE_M=config["GROUP_SIZE_M"], - MUL_ROUTED_WEIGHT=True, - top_k=1, - compute_type_enum=1, - use_fp8_w8a8=True, - use_int8_w8a16=False, - per_channel_quant=True, - even_Ks=moe_intermediate_size % config["BLOCK_SIZE_K"] == 0, - num_warps=config.get("num_warps", 4), - num_stages=config.get("num_stages", 4), + return python_op_wfp8afp8_triton_moe( + x, + getattr(layer, self.added_weight_attrs[0]), + getattr(layer, self.added_scale_attrs[0]), + getattr(layer, self.added_weight_attrs[1]), + getattr(layer, self.added_scale_attrs[1]), + gate_out, + layer.gate_correction_bias, + out, + layer.top_k, + layer.num_local_experts, + layer.moe_intermediate_size, + layer.hidden_size, + layer.topk_method, + getattr(layer, "n_group", 0), + getattr(layer, "topk_group", 0), + getattr(layer, "routed_scaling_factor", 1.0), + getattr(layer, "renormalize", True), + layer.fd_config.scheduler_config.enable_moe_scores_elementwise_fuse, + topk_ids_hookfunc, ) - down_proj_out.reshape_([token_num, top_k, hidden_size]) - out = down_proj_out.sum(axis=1) - - return out - class TensorWiseFP8MoEMethod(QuantMethodBase): """ @@ -1414,6 +1256,241 @@ def python_op_fused_moe_kernel_paddle( return out +def python_op_wfp8afp8_triton_moe_infer_meta( + x, + up_gate_proj_weight, + up_gate_proj_weight_scale, + down_proj_weight, + down_proj_weight_scale, + raw_gate_out, + gate_correction_bias, + out_empty, + top_k: int, + num_local_experts: int, + moe_intermediate_size: int, + hidden_size: int, + topk_method: str, + n_group: int, + topk_group: int, + routed_scaling_factor: float, + renormalize: bool, + enable_moe_scores_elementwise_fuse: bool, + topk_ids_hookfunc, +): + token_num = x.shape[0] + return paddle.static.MetaTensor(shape=[token_num, hidden_size], dtype=x.dtype) + + +@register_custom_python_op( + name="python_op_wfp8afp8_triton_moe", + infer_meta=python_op_wfp8afp8_triton_moe_infer_meta, + input_names=[ + "x", + "up_gate_proj_weight", + "up_gate_proj_weight_scale", + "down_proj_weight", + "down_proj_weight_scale", + "raw_gate_out", + "gate_correction_bias", + "out_empty", + ], + output_names=["out"], + inplace_map={"out_empty": "out"}, +) +def python_op_wfp8afp8_triton_moe( + x: paddle.Tensor, + up_gate_proj_weight: paddle.Tensor, + up_gate_proj_weight_scale: paddle.Tensor, + down_proj_weight: paddle.Tensor, + down_proj_weight_scale: paddle.Tensor, + raw_gate_out: paddle.Tensor, + gate_correction_bias: paddle.Tensor, + out_empty: paddle.Tensor, + top_k: int, + num_local_experts: int, + moe_intermediate_size: int, + hidden_size: int, + topk_method: str, + n_group: int, + topk_group: int, + routed_scaling_factor: float, + renormalize: bool, + enable_moe_scores_elementwise_fuse: bool, + topk_ids_hookfunc, +): + token_num = x.shape[0] + if token_num == 0: + return paddle.zeros([token_num, hidden_size], dtype=x.dtype) + + if topk_method == "noaux_tc": + use_fused = enable_moe_scores_elementwise_fuse and current_platform.is_cuda() + gate_out = raw_gate_out if use_fused else raw_gate_out.cast("float32") + gate_out, topk_weights, topk_ids = get_moe_scores( + gate_out, + n_group, + topk_group, + top_k, + routed_scaling_factor, + gate_correction_bias, + renormalize, + use_fused_cast=use_fused, + ) + else: + gate_out = raw_gate_out.cast("float32") + topk_ids, topk_weights = fastdeploy.model_executor.ops.gpu.moe_topk_select( + gate_out, + gate_correction_bias, + top_k, + True, # apply_norm_weight + False, + ) + + if topk_ids_hookfunc is not None: + topk_ids_hookfunc(topk_ids=topk_ids) + + E, N1, _ = up_gate_proj_weight.shape + + config = { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 32, + "num_warps": 8, + "num_stages": 4, + } + if token_num <= E: + config = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 128, + "GROUP_SIZE_M": 1, + "num_warps": 4, + "num_stages": 4, + } + + from fastdeploy.model_executor.ops.gpu import tritonmoe_preprocess_func + + sorted_token_ids, expert_ids, num_tokens_post_padded = tritonmoe_preprocess_func( + topk_ids, num_local_experts, config["BLOCK_SIZE_M"] + ) + max_possible_num_post_padded = sorted_token_ids.shape[0] + + from .triton_moe_kernels import fused_moe_kernel_paddle + + # --- up-gate proj --- + x_q, x_scale = scaled_fp8_quant(x, use_per_token_if_dynamic=True) + + up_gate_proj_out = paddle.zeros( + [token_num * top_k, moe_intermediate_size * 2], + dtype=x.dtype, + ) + grid_up = ( + ceil_div(max_possible_num_post_padded, config["BLOCK_SIZE_M"]) + * ceil_div(moe_intermediate_size * 2, config["BLOCK_SIZE_N"]), + ) + fused_moe_kernel_paddle[grid_up]( + x_q, + up_gate_proj_weight, + up_gate_proj_out, + x_scale, + up_gate_proj_weight_scale, + None, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + max_possible_num_post_padded, + token_num * top_k, + N=moe_intermediate_size * 2, + K=hidden_size, + stride_am=x_q.strides[0], + stride_ak=x_q.strides[1], + stride_be=up_gate_proj_weight.strides[0], + stride_bk=up_gate_proj_weight.strides[2], + stride_bn=up_gate_proj_weight.strides[1], + stride_cm=up_gate_proj_out.strides[0], + stride_cn=up_gate_proj_out.strides[1], + stride_asm=x_scale.strides[0], + stride_ask=x_scale.strides[1], + stride_bse=up_gate_proj_weight_scale.strides[0], + stride_bsk=up_gate_proj_weight_scale.strides[2], + stride_bsn=up_gate_proj_weight_scale.strides[1], + group_n=-1, + group_k=-1, + BLOCK_SIZE_M=config["BLOCK_SIZE_M"], + BLOCK_SIZE_N=config["BLOCK_SIZE_N"], + BLOCK_SIZE_K=config["BLOCK_SIZE_K"], + GROUP_SIZE_M=config["GROUP_SIZE_M"], + MUL_ROUTED_WEIGHT=False, + top_k=top_k, + compute_type_enum=1, + use_fp8_w8a8=True, + use_int8_w8a16=False, + per_channel_quant=True, + even_Ks=hidden_size % config["BLOCK_SIZE_K"] == 0, + num_warps=config.get("num_warps", 4), + num_stages=config.get("num_stages", 4), + ) + + # --- swiglu + down proj --- + down_proj_input = paddle.incubate.nn.functional.swiglu(up_gate_proj_out) + + x_q_down, x_scale_down = scaled_fp8_quant(down_proj_input, use_per_token_if_dynamic=True) + + down_proj_out = paddle.zeros( + (token_num * top_k, hidden_size), + dtype=x.dtype, + ) + grid_down = ( + ceil_div(max_possible_num_post_padded, config["BLOCK_SIZE_M"]) * ceil_div(hidden_size, config["BLOCK_SIZE_N"]), + ) + fused_moe_kernel_paddle[grid_down]( + x_q_down, + down_proj_weight, + down_proj_out, + x_scale_down, + down_proj_weight_scale, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + max_possible_num_post_padded, + token_num * top_k, + N=hidden_size, + K=moe_intermediate_size, + stride_am=x_q_down.strides[0], + stride_ak=x_q_down.strides[1], + stride_be=down_proj_weight.strides[0], + stride_bk=down_proj_weight.strides[2], + stride_bn=down_proj_weight.strides[1], + stride_cm=down_proj_out.strides[0], + stride_cn=down_proj_out.strides[1], + stride_asm=x_scale_down.strides[0], + stride_ask=x_scale_down.strides[1], + stride_bse=down_proj_weight_scale.strides[0], + stride_bsk=down_proj_weight_scale.strides[2], + stride_bsn=down_proj_weight_scale.strides[1], + group_n=-1, + group_k=-1, + BLOCK_SIZE_M=config["BLOCK_SIZE_M"], + BLOCK_SIZE_N=config["BLOCK_SIZE_N"], + BLOCK_SIZE_K=config["BLOCK_SIZE_K"], + GROUP_SIZE_M=config["GROUP_SIZE_M"], + MUL_ROUTED_WEIGHT=True, + top_k=1, + compute_type_enum=1, + use_fp8_w8a8=True, + use_int8_w8a16=False, + per_channel_quant=True, + even_Ks=moe_intermediate_size % config["BLOCK_SIZE_K"] == 0, + num_warps=config.get("num_warps", 4), + num_stages=config.get("num_stages", 4), + ) + + down_proj_out.reshape_([token_num, top_k, hidden_size]) + out_empty = paddle.sum(down_proj_out, axis=1, out=out_empty) + return out_empty + + class BlockWiseFP8MoEMethod(QuantMethodBase): """ Use Triton Group Gemm to compute Fused BlockWise FP8 Quant MoE. diff --git a/fastdeploy/model_executor/layers/moe/moe.py b/fastdeploy/model_executor/layers/moe/moe.py index 1d4f392208f..a9123a9145a 100644 --- a/fastdeploy/model_executor/layers/moe/moe.py +++ b/fastdeploy/model_executor/layers/moe/moe.py @@ -106,7 +106,6 @@ def get_moe_scores( """ compute moe scores using e_score_correction_bias. """ - assert e_score_correction_bias is not None, "e_score_correction_bias is none!" if envs.FD_USE_PHI_MOE_TOPK: # calculate renormalize and routed_scaling_factor value outside the noaux_tc original_renormalize = renormalize diff --git a/fastdeploy/model_executor/layers/normalization.py b/fastdeploy/model_executor/layers/normalization.py index 7be9e53860b..19df42ac6dc 100644 --- a/fastdeploy/model_executor/layers/normalization.py +++ b/fastdeploy/model_executor/layers/normalization.py @@ -21,9 +21,6 @@ from paddle import nn from fastdeploy.model_executor.forward_meta import ForwardMeta -from fastdeploy.model_executor.graph_optimization.cuda_graph_op import ( - block_wise_cuda_graph_wrap, -) from fastdeploy.platforms import current_platform if current_platform.is_gcu(): @@ -212,7 +209,6 @@ def allgather(self, out, token_num): paddle.distributed.all_gather(multi_outs, out, self.tp_group) return multi_outs[:token_num, :] - @block_wise_cuda_graph_wrap(inputs=["x", "residual_input"], self_attrs=["weight"]) def forward( self, x, @@ -236,24 +232,26 @@ def forward( The `residual_output` is the result of applying the normalization and possibly other operations (like linear transformation) on the `residual_input`. """ + has_residual = residual_input is not None + x_dtype = x.dtype x = x.astype(self.weight.dtype) - if residual_input is not None: + if has_residual: residual_input_dtype = residual_input.dtype residual_input = residual_input.astype(self.weight.dtype) - if residual_input is None: + if not has_residual: residual_out = x use_allreduce_fused = ( self.enable_all_reduce_fusion and self.tp_size > 1 and x.shape[0] <= 2048 - and residual_input is not None + and has_residual and current_platform.is_cuda() ) if proxy_rmsnorm is None: if current_platform.is_gcu(): - if residual_input is None: + if not has_residual: norm_out = rms_norm(x, self.weight, self.eps) return norm_out.astype(x_dtype), residual_out norm_out = self.norm_func(x, residual_input, self.weight, self.eps) @@ -266,7 +264,7 @@ def forward( else: if is_batch_invariant_mode_enabled(): # M-invariant path: per-row Triton kernel, no cross-row reduction - if residual_input is not None: + if has_residual: x = x + residual_input norm_out = rms_norm_batch_invariant(x, self.weight, self.eps), x else: @@ -286,20 +284,16 @@ def forward( else: if use_allreduce_fused: norm_out = flashinfer_allreduce_residual_rmsnorm( - fd_config=self.fd_config, - input_tensor=x, - residual=residual_input, - weight=self.weight, - eps=self.eps, + fd_config=self.fd_config, input_tensor=x, residual=residual_input, weight=self.weight, eps=self.eps ) assert norm_out[0] is not None, "Trtllm-all-reduce fusion failed!" else: - if residual_input is not None: + if has_residual: x = x + residual_input norm_out = proxy_rmsnorm(x, self.weight, self.eps), x out = norm_out[0].astype(x_dtype) - if residual_input is not None: + if has_residual: residual_out = norm_out[1].astype(residual_input_dtype) if self.split_x: diff --git a/fastdeploy/model_executor/layers/quantization/ops/cutlass_scaled_mm.py b/fastdeploy/model_executor/layers/quantization/ops/cutlass_scaled_mm.py index 78321da3381..c26aabdcdd1 100644 --- a/fastdeploy/model_executor/layers/quantization/ops/cutlass_scaled_mm.py +++ b/fastdeploy/model_executor/layers/quantization/ops/cutlass_scaled_mm.py @@ -67,7 +67,12 @@ def cutlass_scaled_mm( ) out = paddle.empty([m, n], dtype=out_dtype) - fastdeploy.model_executor.ops.gpu.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias) + result = fastdeploy.model_executor.ops.gpu.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias) + # In dynamic mode the C++ inplace kernel returns None; use the pre-allocated + # buffer directly. In static (PIR) mode the op returns a new SSA value that + # must be used so downstream ops see the correct dependency edge. + if result is not None: + out = result return out @@ -116,17 +121,23 @@ def scaled_fp8_quant( dynamic_per_token_scaled_fp8_quant, ) - dynamic_per_token_scaled_fp8_quant(output, input, scale, scale_ub) + result = dynamic_per_token_scaled_fp8_quant(output, input, scale, scale_ub) + if result is not None: + output = result else: scale = paddle.zeros([1], dtype=paddle.float32) from fastdeploy.model_executor.ops.gpu import dynamic_scaled_fp8_quant - dynamic_scaled_fp8_quant(output, input, scale) + result = dynamic_scaled_fp8_quant(output, input, scale) + if result is not None: + output = result else: # num_token_padding not implemented for this case # assert (scale.numel() == 1 or num_token_padding is None) from fastdeploy.model_executor.ops.gpu import static_scaled_fp8_quant - static_scaled_fp8_quant(output, input, scale) + result = static_scaled_fp8_quant(output, input, scale) + if result is not None: + output = result return output, scale diff --git a/fastdeploy/model_executor/layers/quantization/ops/scaled_fp8_quant.py b/fastdeploy/model_executor/layers/quantization/ops/scaled_fp8_quant.py index 50c3c6b4347..72ccf09952c 100644 --- a/fastdeploy/model_executor/layers/quantization/ops/scaled_fp8_quant.py +++ b/fastdeploy/model_executor/layers/quantization/ops/scaled_fp8_quant.py @@ -58,22 +58,28 @@ def scaled_fp8_quant( if scale is None: if use_per_token_if_dynamic: - scale = paddle.empty([shape[0], 1], dtype=paddle.float32) + scale = paddle.zeros([shape[0], 1], dtype=paddle.float32) from fastdeploy.model_executor.ops.gpu import ( dynamic_per_token_scaled_fp8_quant, ) - dynamic_per_token_scaled_fp8_quant(output, input, scale, scale_ub) + result = dynamic_per_token_scaled_fp8_quant(output, input, scale, scale_ub) + if result is not None: + output = result else: scale = paddle.zeros([1], dtype=paddle.float32) from fastdeploy.model_executor.ops.gpu import dynamic_scaled_fp8_quant - dynamic_scaled_fp8_quant(output, input, scale) + result = dynamic_scaled_fp8_quant(output, input, scale) + if result is not None: + output = result else: # num_token_padding not implemented for this case # assert (scale.numel() == 1 or num_token_padding is None) from fastdeploy.model_executor.ops.gpu import static_scaled_fp8_quant - static_scaled_fp8_quant(output, input, scale) + result = static_scaled_fp8_quant(output, input, scale) + if result is not None: + output = result return output, scale diff --git a/fastdeploy/worker/gpu_model_runner.py b/fastdeploy/worker/gpu_model_runner.py index 701a7c75fe2..8843e16e729 100644 --- a/fastdeploy/worker/gpu_model_runner.py +++ b/fastdeploy/worker/gpu_model_runner.py @@ -31,6 +31,7 @@ from fastdeploy.engine.pooling_params import PoolingParams from fastdeploy.engine.request import BatchRequest, ImagePosition, Request, RequestType from fastdeploy.model_executor.graph_optimization.utils import ( + prefill_cudagraph_guard, profile_run_guard, sot_warmup_guard, ) @@ -1511,7 +1512,16 @@ def initialize_forward_meta(self, is_dummy_or_profile_run=False): if self.fd_config.parallel_config.enable_chunked_moe: self.forward_meta.max_moe_num_chunk = dist_status.max_moe_num_chunk - only_decode_use_cudagraph = self.use_cudagraph and if_only_decode + # PD prefill workers with piecewise CUDAGraph (graph_opt_level>=1, not full_cuda_graph) only + # capture prefill/mixed graphs (via capture_model_prefill_and_mixed), never decode graphs. + # Exclude such workers from the decode CUDAGraph path to avoid lazy decode capture at runtime. + is_pd_prefill_piecewise = ( + hasattr(self, "graph_opt_config") + and self.fd_config.scheduler_config.splitwise_role == "prefill" + and self.graph_opt_config.graph_opt_level >= 1 + and not self.graph_opt_config.full_cuda_graph + ) + only_decode_use_cudagraph = self.use_cudagraph and if_only_decode and not is_pd_prefill_piecewise # Update config about moe for better performance # TODO(wanglongzhi):Modifying the config at runtime is not appropriate; it needs to be moved to forward_meta. It will be used in MoEMethodBase.apply() @@ -1536,6 +1546,7 @@ def initialize_forward_meta(self, is_dummy_or_profile_run=False): and self.use_cudagraph and self.graph_opt_config.graph_opt_level > 0 and not self.graph_opt_config.full_cuda_graph + and (self.fd_config.scheduler_config.splitwise_role != "prefill" or self.exist_prefill()) ): self.forward_meta.step_use_cudagraph = True @@ -2213,6 +2224,7 @@ def capture_model(self) -> None: time_after_capture = time.perf_counter() logger.info(f"Cuda Graph capturing took {time_after_capture - time_before_capture} seconds") + @prefill_cudagraph_guard(True) @sot_warmup_guard(True) def capture_model_prefill_and_mixed(self) -> None: """ @@ -2222,6 +2234,8 @@ def capture_model_prefill_and_mixed(self) -> None: logger.info("Skipping CUDA graph capture. Please check GraphOptimizationConfig") return time_before_capture = time.perf_counter() + if self.fd_config.parallel_config.use_ep: + self.fd_config.model_config.moe_phase.phase = "prefill" capture_sizes = self.cudagraph_capture_sizes_prefill.copy() for capture_size in sorted(capture_sizes, reverse=True): self._dummy_run( @@ -3144,12 +3158,6 @@ def clear_parameters(self, pid): if self.use_cudagraph: self.model.clear_graph_opt_backend() - if envs.FD_USE_BLOCK_WISE_CUDA_GRAPH: - from fastdeploy.model_executor.graph_optimization.cuda_graph_op import ( - clear_all_block_wise_graphs, - ) - - clear_all_block_wise_graphs() if ( self.speculative_decoding and self.spec_method == SpecMethod.MTP @@ -3259,12 +3267,6 @@ def _clear_cache_for_gdr_weight_update(self): kv_cache_status.value[0] = KVCacheStatus.CLEARING if self.use_cudagraph: self.model.clear_graph_opt_backend() - if envs.FD_USE_BLOCK_WISE_CUDA_GRAPH: - from fastdeploy.model_executor.graph_optimization.cuda_graph_op import ( - clear_all_block_wise_graphs, - ) - - clear_all_block_wise_graphs() if ( self.speculative_decoding and self.spec_method == SpecMethod.MTP @@ -3693,47 +3695,3 @@ def initialize_routing_replay_manager(self): block_table=self.share_inputs["block_tables"], total_block_num=self.num_gpu_blocks, ) - - def capture_block_wise_graphs(self) -> None: - """ - Independent capture loop for block-wise CUDA graphs. - Pre-captures graphs for designated token counts so that at runtime, - matching sizes replay the graph while other sizes fall back to eager. - """ - if not envs.FD_USE_BLOCK_WISE_CUDA_GRAPH: - return - - from fastdeploy.model_executor.graph_optimization.cuda_graph_op import ( - set_block_wise_capturing, - ) - - # Parse capture sizes from env var - sizes_str = envs.FD_BLOCK_WISE_CUDA_GRAPH_SIZES - capture_sizes = sorted([int(s.strip()) for s in sizes_str.split(",") if s.strip()], reverse=True) - if not capture_sizes: - logger.warning("FD_BLOCK_WISE_CUDA_GRAPH_SIZES is empty, skipping block-wise CUDA graph capture") - return - - logger.info(f"Block-wise CUDA graph capture starting for sizes: {sorted(capture_sizes)}") - time_before_capture = time.perf_counter() - - set_block_wise_capturing(True) - try: - for num_tokens in capture_sizes: - batch_size = min(num_tokens, self.scheduler_config.max_num_seqs) - if batch_size < 1: - batch_size = 1 - self._dummy_run( - num_tokens=num_tokens, - batch_size=batch_size, - in_capturing=False, - ) - logger.info(f"Block-wise CUDA graph captured for num_tokens={num_tokens}") - finally: - set_block_wise_capturing(False) - - time_after_capture = time.perf_counter() - logger.info( - f"Block-wise CUDA graph capturing took {time_after_capture - time_before_capture:.3f} seconds " - f"for {len(capture_sizes)} sizes" - ) diff --git a/fastdeploy/worker/gpu_worker.py b/fastdeploy/worker/gpu_worker.py index 9675d086987..118df406d23 100644 --- a/fastdeploy/worker/gpu_worker.py +++ b/fastdeploy/worker/gpu_worker.py @@ -223,30 +223,43 @@ def graph_optimize_and_warm_up_model(self) -> None: """ Perform the warm-up and the graph optimization. - Execution modes: + Execution modes (mixed worker): | Mode | Prefill + Mixed | Decode | |-----------------------------------|--------------------------|--------------------------| | Dynamic (graph_opt_level=0) | Dynamic | Dynamic + CUDAGraph | | Static Full Graph (full=True) | Dynamic | Static + CUDAGraph | | Static Split Graph (full=False) | Static + CUDAGraph | Dynamic + CUDAGraph | + + PD disaggregation: + | Role | graph_opt_level>=1, full=False | Otherwise | + |---------|-------------------------------|------------------------------| + | prefill | Piecewise CUDAGraph (prefill) | Dynamic (cudagraph_only_pref)| + | decode | Dynamic + CUDAGraph | Dynamic + CUDAGraph | """ + splitwise_role = self.fd_config.scheduler_config.splitwise_role + is_pd_prefill = splitwise_role == "prefill" + is_pd_decode = splitwise_role == "decode" + use_piecewise = ( + self.fd_config.graph_opt_config.graph_opt_level >= 1 + and not self.fd_config.graph_opt_config.full_cuda_graph + ) + if self.fd_config.graph_opt_config.graph_opt_level >= 1 and not self.model_runner.use_cudagraph: self.model_runner.sot_warmup() if self.fd_config.graph_opt_config.graph_opt_level >= 1: self.model_runner.vision_encoder_compile() - # Static split graph mode: capture CUDAGraph for prefill/mixed phase - if ( - self.fd_config.graph_opt_config.graph_opt_level >= 1 - and not self.fd_config.graph_opt_config.full_cuda_graph - ): + # Piecewise CUDAGraph capture for prefill/mixed phase. + # In PD disaggregation: only the prefill worker (with piecewise enabled) runs this; + # the decode worker never captures prefill/mixed graphs. + if use_piecewise and not is_pd_decode: self.model_runner.capture_model_prefill_and_mixed() - # Capture CUDAGraph for decode phase (all modes) - self.model_runner.capture_model() - - # Block-wise CUDA graph capture (independent loop) - self.model_runner.capture_block_wise_graphs() + # Decode-phase CUDAGraph capture. + # In PD disaggregation: the prefill worker running piecewise CUDAGraph skips this; + # the decode worker always runs this. + if not (is_pd_prefill and use_piecewise): + self.model_runner.capture_model() # Deterministic mode: reset RNG and share_inputs after warmup. # Warmup _dummy_run() calls consume CUDA RNG state and leave stale diff --git a/tests/graph_optimization/test_block_wise_cuda_graph.py b/tests/graph_optimization/test_block_wise_cuda_graph.py deleted file mode 100644 index 8c39644e3b8..00000000000 --- a/tests/graph_optimization/test_block_wise_cuda_graph.py +++ /dev/null @@ -1,200 +0,0 @@ -# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import os -import shutil -import signal -import subprocess -import sys -import time - -import pytest -import requests - -tests_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) -sys.path.insert(0, tests_dir) - -from e2e.utils.serving_utils import ( - FD_API_PORT, - FD_CACHE_QUEUE_PORT, - FD_ENGINE_QUEUE_PORT, - FD_METRICS_PORT, - PORTS_TO_CLEAN, - clean_ports, - is_port_open, -) - - -@pytest.fixture(scope="session", autouse=True) -def setup_and_run_server(api_url): - """ - Pytest fixture that runs once per test session: - - Cleans ports before tests - - Starts the API server with block-wise CUDA graph env vars enabled - - Waits for server port to open (up to 300 seconds) - - Tears down server after all tests finish - """ - print("Pre-test port cleanup...") - - ports_to_add = [ - FD_API_PORT + 1, - FD_METRICS_PORT + 1, - FD_CACHE_QUEUE_PORT + 1, - FD_ENGINE_QUEUE_PORT + 1, - ] - - for port in ports_to_add: - if port not in PORTS_TO_CLEAN: - PORTS_TO_CLEAN.append(port) - - clean_ports(PORTS_TO_CLEAN) - - print("log dir clean") - if os.path.exists("log") and os.path.isdir("log"): - shutil.rmtree("log") - - base_path = os.getenv("MODEL_PATH") - if base_path: - model_path = os.path.join(base_path, "ernie-4_5-21b-a3b-bf16-paddle") - else: - model_path = "./ernie-4_5-21b-a3b-bf16-paddle" - - log_path = "server.log" - cmd = [ - sys.executable, - "-m", - "fastdeploy.entrypoints.openai.multi_api_server", - "--num-servers", - "2", - "--ports", - f"{FD_API_PORT},{FD_API_PORT + 1}", - "--metrics-ports", - f"{FD_METRICS_PORT},{FD_METRICS_PORT + 1}", - "--args", - "--model", - model_path, - "--engine-worker-queue-port", - f"{FD_ENGINE_QUEUE_PORT},{FD_ENGINE_QUEUE_PORT + 1}", - "--cache-queue-port", - f"{FD_CACHE_QUEUE_PORT},{FD_CACHE_QUEUE_PORT + 1}", - "--tensor-parallel-size", - "1", - "--data-parallel-size", - "2", - "--max-model-len", - "65536", - "--max-num-seqs", - "32", - "--quantization", - "block_wise_fp8", - "--max-num-batched-tokens", - "128", - ] - - # Build env with block-wise CUDA graph enabled - env = os.environ.copy() - env["FD_USE_BLOCK_WISE_CUDA_GRAPH"] = "1" - env["FD_BLOCK_WISE_CUDA_GRAPH_SIZES"] = "128" - env["FD_USE_PHI_FP8_QUANT"] = "0" - env["CUDA_VISIBLE_DEVICES"] = "0,1" - env["FD_BLOCK_WISE_DEBUG"] = "1" - - if os.path.exists("log"): - shutil.rmtree("log") - - with open(log_path, "w") as logfile: - process = subprocess.Popen( - cmd, - stdout=logfile, - stderr=subprocess.STDOUT, - env=env, - start_new_session=True, - ) - - # Wait up to 300 seconds for API server to be ready - for _ in range(300): - if is_port_open("127.0.0.1", FD_API_PORT): - print(f"Server is up on port {FD_API_PORT}") - break - time.sleep(1) - else: - print("[TIMEOUT] API server failed to start in 5 minutes. Cleaning up...") - try: - os.killpg(process.pid, signal.SIGTERM) - except Exception as e: - print(f"Failed to kill process group: {e}") - raise RuntimeError(f"API server did not start on port {FD_API_PORT}") - - yield # Run tests - - print("\n===== Post-test server cleanup... =====") - try: - os.killpg(process.pid, signal.SIGTERM) - time.sleep(10) - print(f"server (pid={process.pid}) terminated") - except Exception as e: - print(f"Failed to terminate API server: {e}") - - clean_ports(PORTS_TO_CLEAN) - - -@pytest.fixture(scope="session") -def api_url(request): - """ - Returns the API endpoint URL for chat completions. - """ - return f"http://0.0.0.0:{FD_API_PORT}/v1/chat/completions" - - -@pytest.fixture -def headers(): - """ - Returns common HTTP request headers. - """ - return {"Content-Type": "application/json"} - - -def send_request(url, payload, timeout=60): - """ - Send a POST request to the specified URL with the given payload. - """ - headers = {"Content-Type": "application/json"} - try: - res = requests.post(url, headers=headers, json=payload, timeout=timeout) - return res - except requests.exceptions.Timeout: - print(f"Request timed out (>{timeout} seconds)") - return None - except requests.exceptions.RequestException as e: - print(f"Request failed: {e}") - return None - - -def test_block_wise_cuda_graph_beijing(api_url): - """ - Verify that block-wise CUDA graph produces correct output. - - With FD_USE_BLOCK_WISE_CUDA_GRAPH=1 and FD_BLOCK_WISE_CUDA_GRAPH_SIZES set, - ask about Tiananmen Square in Beijing and verify the response mentions "北京". - """ - payload = { - "stream": False, - "messages": [{"role": "user", "content": "北京天安门在哪里"}], - "max_tokens": 128, - } - - response = send_request(url=api_url, payload=payload) - print("response: ", response) - assert response is not None, "Request returned None (timeout or connection error)" - assert response.status_code == 200, f"Request failed with status {response.status_code}: {response.text}" diff --git a/tests/graph_optimization/test_cuda_graph_op_unit.py b/tests/graph_optimization/test_cuda_graph_op_unit.py deleted file mode 100644 index ad61d08bfc5..00000000000 --- a/tests/graph_optimization/test_cuda_graph_op_unit.py +++ /dev/null @@ -1,181 +0,0 @@ -# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -""" -In-process unit tests for -``fastdeploy.model_executor.graph_optimization.cuda_graph_op``. - -These tests exercise helper functions and the non-capturing branches of the -``block_wise_cuda_graph_wrap`` decorator to improve coverage. They intentionally -avoid invoking the actual CUDA Graph capture path so that they can run without -spinning up a serving process or recording on a CUDA stream. -""" - -import paddle -import pytest - -import fastdeploy -from fastdeploy.model_executor.graph_optimization import cuda_graph_op - - -class _ToyLayer(paddle.nn.Layer): - """Minimal layer used to exercise the wrap decorator paths.""" - - def __init__(self): - super().__init__() - self.weight = paddle.ones([2, 2], dtype="float32") - - -def _set_env(monkeypatch, **kwargs): - """Helper to override FD env vars dynamically via monkeypatch.""" - for key, value in kwargs.items(): - monkeypatch.setattr(fastdeploy.envs, key, value, raising=False) - - -def test_get_captured_graph_log_returns_copy(monkeypatch): - """get_captured_graph_log should return a list copy of the registry.""" - monkeypatch.setattr(cuda_graph_op, "_CAPTURED_GRAPH_LOG", [("op", ("k",), False)]) - log = cuda_graph_op.get_captured_graph_log() - assert log == [("op", ("k",), False)] - log.append(("other", (), True)) - assert cuda_graph_op._CAPTURED_GRAPH_LOG == [("op", ("k",), False)] - - -def test_dump_captured_graph_summary_paths(monkeypatch): - """Cover early-returns and the summary-print branch of dump_captured_graph_summary.""" - # Debug disabled: early return. - _set_env(monkeypatch, FD_BLOCK_WISE_DEBUG=False) - monkeypatch.setattr(cuda_graph_op, "_CAPTURED_GRAPH_LOG", [("op", (), False)]) - cuda_graph_op.dump_captured_graph_summary() - - # Debug enabled, empty log. - _set_env(monkeypatch, FD_BLOCK_WISE_DEBUG=True) - monkeypatch.setattr(cuda_graph_op, "_CAPTURED_GRAPH_LOG", []) - cuda_graph_op.dump_captured_graph_summary() - - # Debug enabled, non-empty log. - monkeypatch.setattr( - cuda_graph_op, - "_CAPTURED_GRAPH_LOG", - [("a.fwd", ("k1",), True), ("a.fwd", ("k2",), True), ("b.fwd", ("k1",), False)], - ) - cuda_graph_op.dump_captured_graph_summary() - - -def test_clear_all_block_wise_graphs(): - """clear_all_block_wise_graphs should empty every registered shared cache.""" - g, ci, co = {"k": object()}, {"k": object()}, {"k": object()} - snapshot = list(cuda_graph_op._ALL_SHARED_CACHES) - try: - cuda_graph_op._ALL_SHARED_CACHES.clear() - cuda_graph_op._ALL_SHARED_CACHES.append((g, ci, co)) - cuda_graph_op.clear_all_block_wise_graphs() - assert g == {} and ci == {} and co == {} - finally: - cuda_graph_op._ALL_SHARED_CACHES.clear() - cuda_graph_op._ALL_SHARED_CACHES.extend(snapshot) - - -def test_block_wise_wrap_invalid_input_raises(): - """Decorator should raise ValueError when 'inputs' name is not a parameter.""" - with pytest.raises(ValueError): - - @cuda_graph_op.block_wise_cuda_graph_wrap(inputs=["nonexistent"]) - def forward(self, x): - return x - - -def test_block_wise_wrap_disabled_passthrough(monkeypatch): - """When FD_USE_BLOCK_WISE_CUDA_GRAPH is off, wrapper should call eager.""" - _set_env(monkeypatch, FD_USE_BLOCK_WISE_CUDA_GRAPH=False) - - class M(_ToyLayer): - @cuda_graph_op.block_wise_cuda_graph_wrap(inputs=["x"]) - def forward(self, x, residual=None): - return x + 1 - - m = M() - x = paddle.zeros([2, 2], dtype="float32") - out = m.forward(x) - assert paddle.all(out == 1).item() - - -def test_block_wise_wrap_zero_shape_skips(monkeypatch): - """A zero-dim tensor input (positional or keyword) should bypass capture.""" - _set_env(monkeypatch, FD_USE_BLOCK_WISE_CUDA_GRAPH=True) - cuda_graph_op.set_block_wise_capturing(False) - - class M(_ToyLayer): - @cuda_graph_op.block_wise_cuda_graph_wrap(inputs=["x"]) - def forward(self, x, residual=None): - return x - - m = M() - empty = paddle.zeros([0, 4], dtype="float32") - # Positional zero-shape arg: hits the `for a in args` branch. - assert m.forward(empty).shape == [0, 4] - # Keyword zero-shape arg: hits the `for v in kwargs.values()` branch. - assert m.forward(paddle.ones([2, 2]), residual=empty).shape == [2, 2] - - -def test_block_wise_wrap_per_instance_cache_eager_fallback(monkeypatch): - """Per-instance cache init + eager fallback when not in capture phase.""" - _set_env(monkeypatch, FD_USE_BLOCK_WISE_CUDA_GRAPH=True) - cuda_graph_op.set_block_wise_capturing(False) - - class M(_ToyLayer): - @cuda_graph_op.block_wise_cuda_graph_wrap(inputs=["x"]) - def forward(self, x, hook=None): - return x * 2 - - m = M() - x = paddle.ones([2, 2], dtype="float32") - # First call: initializes per-instance cache dicts on self.__dict__. - out1 = m.forward(x, hook=lambda v: v) # callable arg covers callable-key branch - # Second call reuses the existing per-instance cache (try/except hit path). - out2 = m.forward(x) - assert paddle.all(out1 == 2).item() - assert paddle.all(out2 == 2).item() - # The decorator should have stashed cache attributes on the instance. - assert any(name.startswith("_cg_forward_") for name in m.__dict__) - - -def test_block_wise_wrap_custom_key_fn(monkeypatch): - """Custom key_fn path is used to compute the cache key.""" - _set_env(monkeypatch, FD_USE_BLOCK_WISE_CUDA_GRAPH=True) - cuda_graph_op.set_block_wise_capturing(False) - - seen_keys = [] - - def key_fn(x, residual): - k = ("custom", tuple(x.shape) if x is not None else None) - seen_keys.append(k) - return k - - class M(_ToyLayer): - @cuda_graph_op.block_wise_cuda_graph_wrap(inputs=["x"], key_fn=key_fn) - def forward(self, x, residual=None): - return x - - m = M() - m.forward(paddle.ones([3, 3], dtype="float32")) - assert seen_keys and seen_keys[0][0] == "custom" - - -def test_set_block_wise_capturing_toggle(): - """set_block_wise_capturing should mutate the module-level flag.""" - cuda_graph_op.set_block_wise_capturing(True) - assert cuda_graph_op._BLOCK_WISE_CAPTURING is True - cuda_graph_op.set_block_wise_capturing(False) - assert cuda_graph_op._BLOCK_WISE_CAPTURING is False diff --git a/tests/layers/test_kv_cache_int8_dynamic_quant_backend.py b/tests/layers/test_kv_cache_int8_dynamic_quant_backend.py index f679be08b31..ba01ee9ba72 100644 --- a/tests/layers/test_kv_cache_int8_dynamic_quant_backend.py +++ b/tests/layers/test_kv_cache_int8_dynamic_quant_backend.py @@ -103,7 +103,7 @@ def __init__(self): self.graph_opt_config = type( "GraphOptConfig", (), - {"cudagraph_capture_sizes": None}, + {"cudagraph_capture_sizes": None, "full_cuda_graph": True}, )() self.parallel_config = type( "ParallelConfig", diff --git a/tests/model_executor/test_ep.py b/tests/model_executor/test_ep.py index a795249ef55..475d5f5ad33 100644 --- a/tests/model_executor/test_ep.py +++ b/tests/model_executor/test_ep.py @@ -398,9 +398,7 @@ def test_eprunner_moe_select_noaux_tc_without_redundant(monkeypatch): def fake_get_moe_scores(*_args, **_kwargs): return "score", paddle.to_tensor([[0.5]]), paddle.to_tensor([[1]], dtype="int64") - from fastdeploy.model_executor.layers.moe import moe as moe_module - - monkeypatch.setattr(moe_module, "get_moe_scores", fake_get_moe_scores, raising=True) + monkeypatch.setattr(ep, "get_moe_scores", fake_get_moe_scores, raising=True) runner = ep.EPPrefillRunner( top_k=2,