update hybrid kernel to adapt cudagraph

Author: NKNaN

custom_ops/gpu_ops/cpp_extensions.cc

@@ -973,7 +973,8 @@ void HybridMtpNgram(const paddle::Tensor& token_ids_all,
                     const paddle::Tensor& max_dec_len,
                     const int max_ngram_size,
                     const int min_ngram_size,
-                    const int max_draft_tokens);
+                    const int max_draft_tokens,
+                    const bool pad_to_max);
 
 // MTP
 void DraftModelPostprocess(const paddle::Tensor& base_model_draft_tokens,

custom_ops/gpu_ops/speculate_decoding/draft_model/ngram_match_mixed.cu

@@ -144,7 +144,9 @@ __global__ void ngram_match_mixed_gather_kernel(
     int32_t *seq_lens_this_time,
     int64_t draft_tokens_stride,
     int64_t max_batch_size,
-    int threshold) {
+    int threshold,
+    int max_draft_tokens_param,
+    bool pad_to_max) {
   typedef cub::BlockScan<int, NGRAM_GATHER_THREADS> BlockScanInt;
   __shared__ typename BlockScanInt::TempStorage temp_storage1;
   __shared__ typename BlockScanInt::TempStorage temp_storage2;
@@ -203,9 +205,8 @@ __global__ void ngram_match_mixed_gather_kernel(
     }
     actual = min(actual, tentative);
 
-    seq_lens_this_time[tid] = actual;
-
-    // Copy ngram draft tokens from scratch to output
+    // Copy ngram draft tokens from scratch to output FIRST
+    // (so subsequent padding doesn't overwrite real ngram hits)
     int ngram_to_copy = actual - ori;
     if (ngram_to_copy > 0) {
       int64_t *dst = draft_tokens + tid * draft_tokens_stride;
@@ -214,6 +215,37 @@ __global__ void ngram_match_mixed_gather_kernel(
         dst[ori + k] = src[ori + k];
       }
     }
+
+    // === Pad seq_lens_this_time to K+1 for cudagraph stability ===
+    // Hybrid MTP-ngram produces variable seq_lens_this_time depending on how
+    // many ngram positions hit (range: [num_model_steps+1, K+1]). cudagraph
+    // captures launch params (grid dim, kernel args) at capture time; if the
+    // captured slt differs from replay-time slt, downstream kernels read past
+    // valid ranges of cu_seqlens / slot_mapping etc., causing CUDA 700.
+    //
+    // When pad_to_max=true (cudagraph enabled), force slt = K+1 =
+    // max_draft_tokens + 1: positions beyond actual ngram hits get padded
+    // with a placeholder token. The target model will verify these
+    // placeholders and (almost always) reject them, but the verify cost is
+    // fixed per iteration => grid dim is now invariant. When pad_to_max=
+    // false (cudagraph disabled), keep the natural variable slt to avoid
+    // wasting verify compute on placeholders.
+    if (pad_to_max) {
+      int target_slt = max_draft_tokens_param + 1;
+      if (actual < target_slt) {
+        int64_t *dst = draft_tokens + tid * draft_tokens_stride;
+        // Reuse the last valid draft token as placeholder. It is a token the
+        // model could plausibly have produced, so attention math stays
+        // well-defined; rejection happens at the sampler level.
+        int64_t pad_token = (actual > 0) ? dst[actual - 1] : 0;
+        for (int k = actual; k < target_slt; k++) {
+          dst[k] = pad_token;
+        }
+        actual = target_slt;
+      }
+    }
+
+    seq_lens_this_time[tid] = actual;
   }
 }
 
@@ -376,7 +408,8 @@ void HybridMtpNgram(const paddle::Tensor &token_ids_all,
                     const paddle::Tensor &max_dec_len,
                     const int max_ngram_size,
                     const int min_ngram_size,
-                    const int max_draft_tokens) {
+                    const int max_draft_tokens,
+                    const bool pad_to_max) {
   const int64_t max_model_len = token_ids_all.shape()[1];
 
   auto pre_ids_shape = pre_ids.shape();
@@ -404,21 +437,28 @@ void HybridMtpNgram(const paddle::Tensor &token_ids_all,
     // counts tentative > 0, which is equivalent under this invariant.
 
     // Allocate scratch buffers for Phase 1 → Phase 2 communication
+    static paddle::Tensor s_draft_copy_mixed;
+    static paddle::Tensor s_seqlens_copy_mixed;
+    static paddle::Tensor s_seqlens_orig_mixed;
+    static int64_t s_scratch_batch_mixed = 0;
+    static int64_t s_scratch_stride_mixed = 0;
+
+    if (max_batch_size > s_scratch_batch_mixed ||
+        draft_tokens_stride > s_scratch_stride_mixed) {
+      s_draft_copy_mixed = paddle::empty({max_batch_size, draft_tokens_stride},
+                                         paddle::DataType::INT64,
+                                         token_ids_all.place());
+      s_seqlens_copy_mixed = paddle::empty(
+          {max_batch_size}, paddle::DataType::INT32, token_ids_all.place());
+      s_seqlens_orig_mixed = paddle::empty(
+          {max_batch_size}, paddle::DataType::INT32, token_ids_all.place());
+      s_scratch_batch_mixed = max_batch_size;
+      s_scratch_stride_mixed = draft_tokens_stride;
+    }
+    auto &draft_tokens_copy = s_draft_copy_mixed;
+    auto &seq_lens_this_time_copy = s_seqlens_copy_mixed;
+    auto &seq_lens_this_time_orig = s_seqlens_orig_mixed;
 
-    // Scratch copy of draft_tokens (Phase 1 writes tentative tokens here)
-    auto draft_tokens_copy =
-        paddle::empty({max_batch_size, draft_tokens_stride},
-                      paddle::DataType::INT64,
-                      token_ids_all.place());
-
-    // Scratch copy of seq_lens_this_time (Phase 1 writes tentative counts)
-    auto seq_lens_this_time_copy = paddle::empty(
-        {max_batch_size}, paddle::DataType::INT32, token_ids_all.place());
-
-    // Save a copy of original seq_lens_this_time for Phase 2
-    // (Phase 1 reads from the original, Phase 2 needs ori values)
-    auto seq_lens_this_time_orig = paddle::empty(
-        {max_batch_size}, paddle::DataType::INT32, token_ids_all.place());
     cudaMemcpyAsync(seq_lens_this_time_orig.data<int32_t>(),
                     seq_lens_this_time.data<int32_t>(),
                     max_batch_size * sizeof(int32_t),
@@ -462,7 +502,9 @@ void HybridMtpNgram(const paddle::Tensor &token_ids_all,
         const_cast<int32_t *>(seq_lens_this_time.data<int32_t>()),
         draft_tokens_stride,
         max_batch_size,
-        threshold);
+        threshold,
+        max_draft_tokens,
+        pad_to_max);
   } else {
     find_candidate_pred_tokens_mixed(
         token_ids_all.data<int64_t>(),
@@ -496,7 +538,8 @@ PD_BUILD_STATIC_OP(hybrid_mtp_ngram)
              "max_dec_len"})
     .Attrs({"max_ngram_size: int",
             "min_ngram_size: int",
-            "max_draft_tokens: int"})
+            "max_draft_tokens: int",
+            "pad_to_max: bool"})
     .Outputs({"draft_tokens_out", "seq_lens_this_time_out"})
     .SetKernelFn(PD_KERNEL(HybridMtpNgram))
     .SetInplaceMap({{"draft_tokens", "draft_tokens_out"},

fastdeploy/spec_decode/mtp_cuda.py

@@ -393,6 +393,10 @@ def _update_status(self):
             )
 
     def _extend_draft_token_with_ngram_match(self):
+        # pad_to_max forces hybrid kernel to write a fixed seq_lens_this_time
+        # = K + 1, padding unfilled ngram slots with a placeholder draft token.
+        # Required when target cudagraph is enabled (capture-time seq_lens_this_time
+        # must match replay-time seq_lens_this_time).
         hybrid_mtp_ngram(
             self.model_inputs["token_ids_all"],
             self.model_inputs["prompt_lens"],
@@ -406,6 +410,7 @@ def _extend_draft_token_with_ngram_match(self):
             self.max_ngram_size,
             self.min_ngram_size,
             self.max_draft_token_num,
+            self.graph_opt_config.use_cudagraph,
         )
 
     def padding_cudagraph_inputs(self) -> None:

tests/e2e/test_ernie_21b_mtp_ngram.py

@@ -157,9 +157,9 @@ def setup_and_run_server():
         "--cache-queue-port",
         str(FD_CACHE_QUEUE_PORT),
         "--max-model-len",
-        "4096",
+        "32768",
         "--max-num-seqs",
-        "8",
+        "128",
         "--quantization",
         "wint4",
         "--enable-overlap-schedule",

tests/operators/test_hybrid_mtp_ngram.py

@@ -88,6 +88,7 @@ def test_ngram_match_mixed(self):
             self.max_ngram_size,
             self.min_ngram_size,
             self.max_draft_tokens,
+            False,
         )
 
         np.testing.assert_allclose(self.seq_lens_this_time.numpy(), self.ref_seq_lens_this_time)

tests/spec_decode/test_ngram_gpu_kernel.py

@@ -846,6 +846,7 @@ def test_correctness_basic(self):
             max_ngram_size,
             min_ngram_size,
             max_draft_tokens,
+            False,
         )
         paddle.device.synchronize()
 
@@ -889,6 +890,7 @@ def test_correctness_varied_seeds(self):
                     3,
                     1,
                     10,
+                    False,
                 )
                 paddle.device.synchronize()
                 np.testing.assert_array_equal(gpu_data["seq_lens_this_time"].numpy(), cpu_slt)
@@ -936,6 +938,7 @@ def test_large_batch_long_seq(self):
                 3,
                 1,
                 10,
+                False,
             )
             paddle.device.synchronize()
         finally:
@@ -979,6 +982,7 @@ def test_single_batch_long_seq(self):
             3,
             1,
             10,
+            False,
         )
         paddle.device.synchronize()
         np.testing.assert_array_equal(gpu_data["seq_lens_this_time"].numpy(), cpu_slt)
@@ -1022,6 +1026,7 @@ def test_many_short_seqs(self):
                 3,
                 1,
                 10,
+                False,
             )
             paddle.device.synchronize()
         finally:

tests/worker/test_reorder_split_prefill_and_decode.py

@@ -1,4 +1,4 @@
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import paddle
 import pytest
@@ -12,7 +12,11 @@
     SpeculativeConfig,
     StructuredOutputsConfig,
 )
-from fastdeploy.worker.input_batch import InputBatch, reorder_split_prefill_and_decode
+from fastdeploy.worker.input_batch import (
+    InputBatch,
+    ProposerInputBatch,
+    reorder_split_prefill_and_decode,
+)
 
 
 def create_mock_config():
@@ -61,6 +65,7 @@ def create_mock_config():
     scheduler_config = Mock(spec=SchedulerConfig)
     scheduler_config.max_num_seqs = 10
     scheduler_config.max_num_batched_tokens = 2048
+    scheduler_config.max_extra_num_batched_tokens = 0
 
     speculative_config = Mock(spec=SpeculativeConfig)
     speculative_config.method = None
@@ -315,5 +320,119 @@ def test_reorder_all_prefill(self):
             assert paddle.equal_all(input_batch.input_ids[i], original_data[i])
 
 
+class TestProposerInputBatchReset:
+    """Cover ProposerInputBatch.reset_model_inputs CUDA + token_ids_all branch
+    (fastdeploy/worker/input_batch.py:972-985)."""
+
+    def _make_config(self):
+        # Enable spec_decoding path so InputBatch.init_share_inputs allocates
+        # cu_seqlens_q_output / draft_tokens / accept_num etc.
+        fd_config = create_mock_config()
+        fd_config.speculative_config.method = "mtp"
+        fd_config.speculative_config.num_speculative_tokens = 1
+        fd_config.speculative_config.num_model_steps = 1
+        return fd_config
+
+    def _make_target(self, fd_config):
+        target = InputBatch(fd_config)
+        target.init_share_inputs()
+        return target
+
+    def _make_proposer(self, fd_config, target):
+        """Construct a ProposerInputBatch and manually populate only the
+        attributes that `reset_model_inputs` writes via `fill_paddle_tensor`.
+        Skipping full `init_share_inputs` avoids depending on rope_emb,
+        attention backends, and other heavy setup unrelated to the branch
+        under test."""
+        proposer = ProposerInputBatch(fd_config, target)
+
+        max_num_seqs = fd_config.scheduler_config.max_num_seqs
+        hidden_size = fd_config.model_config.hidden_size
+        max_draft_token_num = fd_config.speculative_config.num_speculative_tokens
+
+        proposer.target_hidden_states = paddle.full([max_num_seqs, hidden_size], 0, dtype="bfloat16")
+        proposer.draft_tokens = paddle.full([max_num_seqs, max_draft_token_num + 1], -1, dtype="int64")
+        proposer.is_block_step = paddle.full([max_num_seqs, 1], False, dtype="bool")
+        proposer.batch_drop = paddle.full([max_num_seqs, 1], False, dtype="bool")
+        proposer.used_list_len = paddle.full([max_num_seqs], 0, dtype="int32")
+        proposer.first_token_hidden_states = paddle.full([max_num_seqs, hidden_size], -1)
+        proposer.batch_token_num = paddle.full([max_num_seqs], 0, dtype="int32")
+        proposer.next_token_num = paddle.full([max_num_seqs], 0, dtype="int32")
+        proposer.cu_batch_token_offset = paddle.full([max_num_seqs + 1], 0, dtype="int32")
+        proposer.cu_next_token_offset = paddle.full([max_num_seqs + 1], 0, dtype="int32")
+        proposer.mask_rollback = paddle.full([max_num_seqs, 1], 0, dtype="int32")
+        proposer.recompute_token_num = paddle.full([max_num_seqs, 1], 0, dtype="int32")
+        return proposer
+
+    @patch("fastdeploy.worker.input_batch.current_platform")
+    def test_reset_rebinds_token_ids_all_on_cuda(self, mock_platform):
+        """When current_platform.is_cuda() and target has token_ids_all,
+        reset_model_inputs must re-pull token_ids_all from target (line 973)
+        and rebuild pre_ids from target.token_ids_all[bs_idx, prompt_len:]."""
+        mock_platform.is_cuda.return_value = True
+        mock_platform.is_xpu.return_value = False
+
+        fd_config = self._make_config()
+        target = self._make_target(fd_config)
+        proposer = self._make_proposer(fd_config, target)
+
+        max_num_seqs = fd_config.scheduler_config.max_num_seqs
+        max_model_len = fd_config.model_config.max_model_len
+
+        # Rebind target.token_ids_all to a NEW tensor with known content so
+        # we can distinguish "reset re-pulled it" from "init_share_inputs
+        # already bound to the old reference".
+        new_token_ids_all = paddle.arange(max_num_seqs * max_model_len, dtype="int64").reshape(
+            [max_num_seqs, max_model_len]
+        )
+        target.token_ids_all = new_token_ids_all
+
+        # Set a non-zero prompt_len so the slice [:, prompt_len:] is non-trivial.
+        prompt_len_value = 3
+        target.prompt_lens = paddle.full([max_num_seqs, 1], prompt_len_value, dtype="int64")
+
+        proposer.reset_model_inputs()
+
+        # Line 973 effect: token_ids_all rebound to target's new tensor.
+        assert proposer.token_ids_all is new_token_ids_all
+
+        # Line 975-985 effect: pre_ids has correct shape and the prefix is
+        # token_ids_all[bs_idx, prompt_len:], suffix remains -1.
+        assert proposer.pre_ids.shape == [max_num_seqs, max_model_len]
+        valid_len = max_model_len - prompt_len_value
+        expected_prefix = new_token_ids_all[:, prompt_len_value:]
+        assert paddle.equal_all(proposer.pre_ids[:, :valid_len], expected_prefix)
+        assert paddle.equal_all(
+            proposer.pre_ids[:, valid_len:],
+            paddle.full([max_num_seqs, prompt_len_value], -1, dtype="int64"),
+        )
+
+    @patch("fastdeploy.worker.input_batch.current_platform")
+    def test_reset_falls_back_to_pre_ids_clone_when_no_token_ids_all(self, mock_platform):
+        """When current_platform.is_cuda() but target lacks token_ids_all,
+        reset_model_inputs takes the else branch (line 986-988): clone pre_ids,
+        set token_ids_all to None."""
+        mock_platform.is_cuda.return_value = True
+        mock_platform.is_xpu.return_value = False
+
+        fd_config = self._make_config()
+        target = self._make_target(fd_config)
+        proposer = self._make_proposer(fd_config, target)
+
+        # Remove token_ids_all from target so the else branch fires.
+        del target.token_ids_all
+        # Provide a recognizable pre_ids on target.
+        max_num_seqs = fd_config.scheduler_config.max_num_seqs
+        max_model_len = fd_config.model_config.max_model_len
+        target.pre_ids = paddle.full([max_num_seqs, max_model_len], 42, dtype="int64")
+
+        proposer.reset_model_inputs()
+
+        assert proposer.token_ids_all is None
+        assert paddle.equal_all(proposer.pre_ids, paddle.full([max_num_seqs, max_model_len], 42, dtype="int64"))
+        # Clone, not reference share.
+        assert proposer.pre_ids is not target.pre_ids
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])