bugfix: fix mtp sampling mixed issue. (#1128)

Author: phantomlei3

xllm/core/framework/sampling/rejection_sampler.cpp

@@ -105,8 +105,11 @@ RejectionSampler::RejectionSampler(
       rate_controller_(rate_controller),
       enable_fused_kernel_(enable_fused_kernel) {
   CHECK(do_sample.defined());
-  // [batch_size, 1]
-  do_sample_ = do_sample.unsqueeze_(/*dim=*/-1);
+  // Keep a private expanded view and do not mutate the caller-owned tensor.
+  // The same SamplingParameters object is reused later by MTP draft extend.
+  // An in-place unsqueeze here corrupts Sampler::forward() mixed-mode shape
+  // assumptions and can broadcast sampled token ids into 2D.
+  do_sample_ = do_sample.unsqueeze(/*dim=*/-1);
 }
 
 // draft_token_ids: [batch_size, n_speculative_tokens]

xllm/core/framework/sampling/rejection_sampler_test.cpp

@@ -254,6 +254,55 @@ TEST(RejectionSamplerTest, LogProbs) {
   EXPECT_TRUE(torch::equal(output.top_tokens, top_k_indices));
 }
 
+TEST(RejectionSamplerTest, ConstructorDoesNotMutateDoSampleShape) {
+  const auto device = get_test_device();
+  auto do_sample = torch::tensor({false, true}, torch::device(device));
+
+  ASSERT_EQ(do_sample.dim(), 1);
+  ASSERT_EQ(do_sample.sizes(), torch::IntArrayRef({2}));
+
+  RejectionSampler sampler(do_sample,
+                           do_sample.all().item<bool>(),
+                           !do_sample.any().item<bool>(),
+                           /*logprobs=*/false,
+                           /*max_top_logprobs=*/0);
+
+  EXPECT_EQ(do_sample.dim(), 1);
+  EXPECT_EQ(do_sample.sizes(), torch::IntArrayRef({2}));
+}
+
+TEST(RejectionSamplerTest,
+     ReusingDoSampleAfterRejectionSamplerKeepsSamplerOutput1D) {
+  const auto options = get_test_options(torch::kFloat32);
+  const auto device = get_test_device();
+  auto do_sample = torch::tensor({false, true}, torch::device(device));
+
+  RejectionSampler rejection_sampler(do_sample,
+                                     do_sample.all().item<bool>(),
+                                     !do_sample.any().item<bool>(),
+                                     /*logprobs=*/false,
+                                     /*max_top_logprobs=*/0);
+  (void)rejection_sampler;
+
+  SamplingParameters params;
+  params.selected_token_idxes =
+      torch::tensor({0, 1}, torch::dtype(torch::kInt64).device(device));
+  params.sample_idxes =
+      torch::tensor({0, 1}, torch::dtype(torch::kInt64).device(device));
+  params.do_sample = do_sample;
+  params.all_random_sample = false;
+  params.all_greedy_sample = false;
+
+  auto logits =
+      torch::tensor({{3.0f, 1.0f, 0.5f}, {0.1f, 0.2f, 4.0f}}, options);
+  auto output = Sampler().forward(logits, params);
+
+  EXPECT_EQ(output.probs.dim(), 2);
+  EXPECT_EQ(output.probs.size(0), 2);
+  EXPECT_EQ(output.next_tokens.dim(), 1);
+  EXPECT_EQ(output.next_tokens.size(0), 2);
+}
+
 TEST(RejectionSamplerTest, Random) {
   const auto options = get_test_options(torch::kFloat32);
 

xllm/core/framework/sampling/sampler.cpp

@@ -51,6 +51,10 @@ SampleOutput Sampler::forward(torch::Tensor& logits,
     sample_logits = logits.index_select(/*dim=*/0, params.sample_idxes);
   }
 
+  CHECK(params.do_sample.defined()) << "params.do_sample must be defined";
+  CHECK_EQ(params.do_sample.dim(), 1)
+      << "params.do_sample must be 1D [num_seqs], got "
+      << params.do_sample.sizes();
   // same batch size
   CHECK_EQ(sample_logits.size(0), params.do_sample.size(0));
 

xllm/core/runtime/mtp_worker_impl.cpp

@@ -507,6 +507,16 @@ std::optional<ForwardOutput> MTPWorkerImpl::run_validate(
 
 void MTPWorkerImpl::process_draft_sample_output(SampleOutput& sample_output) {
   if (sample_output.probs.defined()) {
+    CHECK(sample_output.next_tokens.defined())
+        << "draft sample_output.next_tokens must be defined when probs exist";
+    CHECK_EQ(sample_output.next_tokens.dim(), 1)
+        << "MTP draft cache expects next_tokens [batch], got "
+        << sample_output.next_tokens.sizes();
+    CHECK(sample_output.probs.dim() == 1 || sample_output.probs.dim() == 2)
+        << "MTP draft cache expects probs [batch] or [batch,vocab], got "
+        << sample_output.probs.sizes();
+    CHECK_EQ(sample_output.probs.size(0), sample_output.next_tokens.size(0))
+        << "MTP draft cache probs/token batch mismatch";
     // Cache always stores selected-only draft probs [batch_size] to reduce HBM.
     sample_output.probs = specBuilder::draftProbs::compress_for_cache(
         sample_output.probs, sample_output.next_tokens);