[BugFix][Speculative Decoding] Fix MTP triton sampler seed (#7946)

* add FD_ENABLE_TOP_P_ONE_OPT=0

* mtp support triton seed

* fix mtp triton sampler

* check test

Author: ckl117

fastdeploy/model_executor/layers/sample/ops/top_k_top_p_triton.py

@@ -947,7 +947,6 @@ def _seeded_gumbel_kernel(
     """Generate -log(u) with per-row Philox seeds, fully on GPU."""
     pid = tl.program_id(0)
     seed = tl.load(SEEDS_ptr + pid)
-    seed = seed.to(tl.int32)
     for start in tl.range(0, VOCAB_SIZE, BLOCK_SIZE):
         offsets = start + tl.arange(0, BLOCK_SIZE)
         mask = offsets < VOCAB_SIZE

fastdeploy/model_executor/layers/sample/sampler.py

@@ -927,6 +927,7 @@ def _verify_and_sample(
         increment_value: int,
         accept_all_drafts: bool = False,
         reject_all_drafts: bool = False,
+        topp_seed: Optional[paddle.Tensor] = None,
     ) -> SamplerOutput:
         """
         Verify draft tokens against target model output and produce final samples.
@@ -959,7 +960,7 @@ def _verify_and_sample(
 
         if self.verify_strategy == VerifyStrategy.TARGET_MATCH:
             if FD_SAMPLING_CLASS.lower() == "triton":
-                target_tokens = _random_sample(probs, topp_seed=sampling_metadata.seed)
+                target_tokens = _random_sample(probs, topp_seed=topp_seed)
             else:
                 # Only TARGET_MATCH needs stochastic sampling
                 top_p, top_k, topp_seed = build_sampling_params(
@@ -1038,6 +1039,7 @@ def _normal_sample(
         probs: paddle.Tensor,
         sampling_metadata: SamplingMetadata,
         share_inputs: List[paddle.Tensor],
+        topp_seed: Optional[paddle.Tensor],
     ) -> SamplerOutput:
         """
         Normal sampling without draft token verification.
@@ -1060,7 +1062,7 @@ def _normal_sample(
 
         # Sample tokens
         if FD_SAMPLING_CLASS.lower() == "triton":
-            next_tokens = _random_sample(probs, topp_seed=sampling_metadata.seed)
+            next_tokens = _random_sample(probs, topp_seed=topp_seed)
         else:
             next_tokens = _sample_from_probs(
                 probs,
@@ -1164,6 +1166,7 @@ def forward_cuda(
             )
 
         logits_ori = None
+        topp_seed = None
         if FD_SAMPLING_CLASS.lower() == "triton":
             logits_ori = logits.clone()
             top_p, top_k, _ = build_sampling_params(
@@ -1187,7 +1190,7 @@ def forward_cuda(
         # Route based on spec_method
         is_naive = self.spec_method is None or self.spec_method == SpecMethod.NAIVE
         if is_naive:
-            sampler_output = self._normal_sample(logits, probs, sampling_metadata, share_inputs)
+            sampler_output = self._normal_sample(logits, probs, sampling_metadata, share_inputs, topp_seed=topp_seed)
         else:
             sampler_output = self._verify_and_sample(
                 logits,
@@ -1199,6 +1202,7 @@ def forward_cuda(
                 increment_value,
                 accept_all_drafts,
                 reject_all_drafts,
+                topp_seed=topp_seed,
             )
 
         # Build logprobs via unified path (outside of sampling logic)

tests/layers/test_triton_sampler.py

@@ -121,6 +121,7 @@ def triton_mode(monkeypatch):
     import fastdeploy.envs as envs
 
     monkeypatch.setattr(envs, "FD_SAMPLING_CLASS", "triton")
+    monkeypatch.setattr("fastdeploy.model_executor.layers.sample.sampler.FD_SAMPLING_CLASS", "triton")
 
 
 def _create_metadata(batch_size=1, min_seq_len=1, max_seq_len=3, max_num_logprobs=None, **overrides):
@@ -343,6 +344,7 @@ def test_verify_and_sample_target_match_triton(self, mock_ops, triton_mode, monk
         m = _create_metadata(batch_size=1)
         logits = paddle.to_tensor([[1.0, 2.0, 3.0]], dtype="float32")
         probs = paddle.nn.functional.softmax(logits, axis=-1)
+        seeds = paddle.ones([probs.shape[0], 1], dtype="int64")
 
         out = sampler._verify_and_sample(
             logits,
@@ -352,6 +354,7 @@ def test_verify_and_sample_target_match_triton(self, mock_ops, triton_mode, monk
             share_inputs=_spec_share_inputs(),
             token_num_output_cpu=1,
             increment_value=1,
+            topp_seed=seeds,
         )
         assert out.sampled_token_ids is not None
 
@@ -366,8 +369,9 @@ def test_normal_sample_triton(self, mock_ops, triton_mode, monkeypatch):
         m = _create_metadata(batch_size=1)
         logits = paddle.to_tensor([[1.0, 2.0, 3.0]], dtype="float32")
         probs = paddle.nn.functional.softmax(logits, axis=-1)
+        seeds = paddle.ones([probs.shape[0], 1], dtype="int64")
 
-        out = sampler._normal_sample(logits, probs, m, share_inputs=_spec_share_inputs())
+        out = sampler._normal_sample(logits, probs, m, share_inputs=_spec_share_inputs(), topp_seed=seeds)
         assert out.sampled_token_ids is not None
 
     def test_forward_cuda_triton_logit_mask(self, mock_ops, triton_mode, monkeypatch):