Revert blockwise CUDAGraph and support piecewise CUDAGraph in prefill

Author: ZhangX-21

fastdeploy/config.py

@@ -1211,6 +1211,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)

fastdeploy/envs.py

@@ -263,16 +263,6 @@ def _validate_split_kv_size(value: int) -> int:
     "FD_SAVE_OUTPUT_CACHE_FOR_PREEMPTED_REQUEST": lambda: bool(
         int(os.getenv("FD_SAVE_OUTPUT_CACHE_FOR_PREEMPTED_REQUEST", "1"))
     ),
-    # 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

fastdeploy/model_executor/graph_optimization/cuda_graph_op.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)

fastdeploy/model_executor/graph_optimization/utils.py

@@ -321,17 +321,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,

fastdeploy/model_executor/layers/attention/append_attn_backend.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.