Fix draft logprobs zeros bug and add acceptance sanity checks

Bug #1: EAGLE tree proposal returned zeros for draft_logprobs
- Root cause: When using topk for tree branching, code set draft_logp_list=None,
  then created zeros tensor as fallback (lines 850-851)
- Fix: Compute actual log-probs from logits using log_softmax + gather
- Applied at 2 locations: root level (lines 698-704) and tree levels (lines 839-846)

Bug #2: Added diagnostic logging in rejection sampler
- Log draft_p (nonzero) min/med/max to detect zeros
- Log p_target min/med/max to detect degenerate softmax
- Helps identify if target logits are masked/filtered before sampling

Expected results after fix:
- draft_logp: -3.2/-1.6/-0.0 (real log-probs, all ≤ 0) instead of 0/0/0
- p_target: 1e-6/1e-3/0.7 (realistic distribution) instead of 1/1/1
- Acceptance rate: 30-70% instead of 0%

Files changed:
- vllm/v1/spec_decode/eagle.py: Fix draft_logp computation
- vllm/v1/sample/rejection_sampler.py: Add sanity logging

Author: yuz207

vllm/v1/sample/rejection_sampler.py

@@ -124,6 +124,26 @@ def forward(
             sampling_metadata,
         )
 
+        # Sanity checks: Inspect raw values BEFORE any processing
+        if draft_probs is not None:
+            draft_p_nonzero = draft_probs[draft_probs > 0]
+            if draft_p_nonzero.numel() > 0:
+                print(f"[SANITY] draft_p (nonzero) min/med/max: "
+                      f"{draft_p_nonzero.min():.3e}/"
+                      f"{draft_p_nonzero.median():.3e}/"
+                      f"{draft_p_nonzero.max():.3e}",
+                      file=sys.stderr, flush=True)
+            else:
+                print(f"[SANITY] draft_p: ALL ZEROS!", file=sys.stderr, flush=True)
+
+        # Check target probabilities for the chosen draft tokens
+        target_p_check = target_probs.gather(-1, metadata.draft_token_ids.unsqueeze(-1)).squeeze(-1)
+        print(f"[SANITY] p_target min/med/max: "
+              f"{target_p_check.min():.3e}/"
+              f"{target_p_check.median():.3e}/"
+              f"{target_p_check.max():.3e}",
+              file=sys.stderr, flush=True)
+
         output_token_ids = rejection_sample(
             metadata.draft_token_ids,
             metadata.num_draft_tokens,

vllm/v1/spec_decode/eagle.py

@@ -695,10 +695,13 @@ def propose_tree(
             draft_token_ids_list = [draft_token_ids]
             draft_logp_list = [draft_logp]
         else:
+            # Top-k branching: compute actual log-probs for chosen tokens
             draft_token_ids = torch.topk(logits, num_children,
                                          dim=-1).indices.view(batch_size, -1)
+            log_probs = torch.log_softmax(logits, dim=-1)
+            draft_logp = log_probs.gather(-1, draft_token_ids).view(batch_size, -1)
             draft_token_ids_list = [draft_token_ids]
-            draft_logp_list = None  # No valid draft probs for deterministic top-k
+            draft_logp_list = [draft_logp]
         draft_hidden_states = hidden_states.view(batch_size, 1, -1)
 
         # Initialize empty tensors for concatenation with the level outputs.
@@ -833,23 +836,22 @@ def propose_tree(
                 draft_token_ids_list.append(draft_token_ids)
                 draft_logp_list.append(draft_logp)
             else:
+                # Top-k branching: compute actual log-probs for chosen tokens
                 draft_token_ids = torch.topk(logits, num_children,
                                              dim=-1).indices.view(
                                                  batch_size, -1)
+                log_probs = torch.log_softmax(logits, dim=-1)
+                draft_logp = log_probs.gather(-1, draft_token_ids).view(batch_size, -1)
                 draft_token_ids_list.append(draft_token_ids)
-                if draft_logp_list is not None:
-                    draft_logp_list = None  # Mixed modes -> no valid logprobs
+                draft_logp_list.append(draft_logp)
 
             # Update the # drafts counters for the next tree level.
             level_num_drafts = self.cu_drafts_per_level[level +
                                                         1] - total_num_drafts
             total_num_drafts = self.cu_drafts_per_level[level + 1]
 
         # Return both token IDs and logprobs
-        # If draft_logp_list is None (top-k was used), return zeros
-        if draft_logp_list is None:
-            draft_logp_list = [torch.zeros_like(draft_token_ids_list[0], dtype=torch.float32)]
-
+        # All branches now compute real logprobs, no need for fallback
         return draft_token_ids_list, draft_logp_list
 
     def prepare_inputs(