[Spec] Support quantized target lm_head in the DFlash2 selector (#35496)

Co-authored-by: LING ZHI <1747985437lz@gmail.com>

python/sglang/srt/models/dflash.py

@@ -23,7 +23,10 @@ from sglang.srt.layers.linear import (
     QKVParallelLinear,
     RowParallelLinear,
 )
-from sglang.srt.layers.logits_processor import LogitsProcessorOutput
+from sglang.srt.layers.logits_processor import (
+    LogitsProcessorOutput,
+    should_apply_lm_head_quant_method,
+)
 from sglang.srt.layers.radix_attention import AttentionType, RadixAttention
 from sglang.srt.layers.rotary_embedding import get_rope
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -59,6 +62,22 @@ def _radix_topk(scores: torch.Tensor, k: int) -> Tuple[torch.Tensor, torch.Tenso
     return torch.topk(scores, k, dim=-1)
 
 
+def _project_candidate_logits(
+    hidden: torch.Tensor, lm_head: nn.Module, *, num_org: int, use_quant_head: bool
+) -> torch.Tensor:
+    """Project draft hiddens through the target head, restricted to the org vocab."""
+    if not use_quant_head:
+        weight = lm_head.weight
+        return torch.matmul(hidden.to(weight.dtype), weight[:num_org].T)
+    # A packed weight can't be row-sliced to the org vocab like the dense path,
+    # and flashinfer's radix top-k rejects the crop view (non-contiguous), so
+    # mask the padded tail out of the top-k instead.
+    logits = lm_head.quant_method.apply(lm_head, hidden, None).contiguous()
+    if logits.shape[-1] > num_org:
+        logits[:, num_org:] = float("-inf")
+    return logits
+
+
 def _get_dflash_attention_type(config, *, default: AttentionType) -> AttentionType:
     """Honor explicit causality while preserving legacy layer defaults."""
     text_config = config.get_text_config()
@@ -965,18 +984,31 @@ class DFlash2DraftModel(DFlashDraftModel):
         # The worker screens the head before capture, but its eager fallback
         # (_propose_selector_block) attaches whatever the target has.
         weight = getattr(self.lm_head, "weight", None)
-        if not is_dense_head_weight(weight):
+        quant_method = getattr(self.lm_head, "quant_method", None)
+        use_quant_head = should_apply_lm_head_quant_method(self.lm_head, quant_method)
+        if not use_quant_head and not is_dense_head_weight(weight):
             raise RuntimeError(
-                "DFlash2 selector requires a dense FP16/BF16/FP32 target lm_head."
+                "DFlash2 selector requires a dense FP16/BF16/FP32 target lm_head "
+                "or a supported lm_head.quant_method."
             )
-        hidden = hidden.to(weight.dtype)
         if get_parallel().tp_size == 1:
             org = int(self.lm_head.org_vocab_size)
-            vals, ids = _radix_topk(torch.matmul(hidden, weight[:org].T), k)
+            vals, ids = _radix_topk(
+                _project_candidate_logits(
+                    hidden, self.lm_head, num_org=org, use_quant_head=use_quant_head
+                ),
+                k,
+            )
             return ids.long(), self._transform_unary_logits(vals)
         shard = self.lm_head.shard_indices
         vals, ids = _radix_topk(
-            torch.matmul(hidden, weight[: int(shard.num_org_elements)].T), k
+            _project_candidate_logits(
+                hidden,
+                self.lm_head,
+                num_org=int(shard.num_org_elements),
+                use_quant_head=use_quant_head,
+            ),
+            k,
         )
         global_ids = ids.long() + int(shard.org_vocab_start_index)
         gathered_vals = tensor_model_parallel_all_gather(vals.float(), dim=-1)

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -20,6 +20,7 @@ from sglang.srt.configs.hybrid_arch import mambaish_config
 from sglang.srt.distributed import get_tp_group
 from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.environ import envs
+from sglang.srt.layers.logits_processor import should_apply_lm_head_quant_method
 from sglang.srt.layers.logprob_processor import compute_spec_logprobs
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
@@ -482,14 +483,17 @@ class DFlashWorkerV2(BaseSpecWorker):
         lm_head = getattr(target_model, "lm_head", None)
         if lm_head is None:
             return _eager("no target lm_head")
-        if not hasattr(lm_head, "weight"):
-            return _eager("quantized lm_head has no dense weight")
-        if not is_dense_head_weight(lm_head.weight):
-            # Quantized lm_head (FP8/INT) would break the static matmul.
-            return _eager("quantized lm_head")
 
         if self.selector is not None:
             # compute_candidates needs the target lm_head attached before capture.
+            # A gate-admitted quantized head is capture-safe: the target's own
+            # logits path already runs the same kernel under CUDA graphs.
+            if not is_dense_head_weight(
+                getattr(lm_head, "weight", None)
+            ) and not should_apply_lm_head_quant_method(
+                lm_head, getattr(lm_head, "quant_method", None)
+            ):
+                return _eager("unsupported quantized lm_head")
             self.draft_model.lm_head = lm_head
             if self.ps.tp_rank == 0:
                 logger.info(
@@ -502,6 +506,11 @@ class DFlashWorkerV2(BaseSpecWorker):
                 max_bs=max(get_exec().graph.cuda_graph_config.decode.bs),
                 device=self.device,
             )
+        if not hasattr(lm_head, "weight"):
+            return _eager("quantized lm_head has no dense weight")
+        if not is_dense_head_weight(lm_head.weight):
+            # Quantized lm_head (FP8/INT) would break the static matmul.
+            return _eager("quantized lm_head")
         tp_group = get_tp_group()
         if not hasattr(lm_head, "shard_indices"):
             if tp_group.world_size != 1: