[Graph Optimization]Revert blockwise CUDAGraph and support piecewise CUDAGraph in prefill

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，所以必须！

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，所以必须！

custom_ops/gpu_ops/grouped_topk_kernels.cu

@@ -767,6 +767,8 @@ std::vector<paddle::DataType> GroupedTopkInferDtype(
 std::vector<std::vector<int64_t>> GroupedTopkInferShape(
     const std::vector<int64_t>& gating_output_shape,
     const std::vector<int64_t>&,
+    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];

custom_ops/gpu_ops/moe/tritonmoe_preprocess.cu

@@ -29,10 +29,25 @@ std::vector<std::vector<int64_t>> tritonmoe_preprocessInferShape(
     const std::vector<int64_t>& 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<int>(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;

custom_ops/gpu_ops/noaux_tc.cu

@@ -67,6 +67,8 @@ std::vector<paddle::DataType> NoauxTcInferDtype(
 std::vector<std::vector<int64_t>> NoauxTcInferShape(
     const std::vector<int64_t>& scores_shape,
     const std::vector<int64_t>&,
+    const int n_group,
+    const int topk_group,
     const int topk) {
   auto num_tokens = scores_shape[0];
   auto topk_values_shape = std::vector<int64_t>{num_tokens, topk};

custom_ops/gpu_ops/noaux_tc_redundant.cu

@@ -74,6 +74,11 @@ std::vector<paddle::DataType> NoauxTcRedundantInferDtype(
 std::vector<std::vector<int64_t>> NoauxTcRedundantInferShape(
     const std::vector<int64_t>& scores_shape,
     const std::vector<int64_t>&,
+    const std::vector<int64_t>&,
+    const std::vector<int64_t>&,
+    const std::vector<int64_t>&,
+    const int n_group,
+    const int topk_group,
     const int topk) {
   auto num_tokens = scores_shape[0];
   auto topk_values_shape = std::vector<int64_t>{num_tokens, topk};

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(

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