[Spec] DFlash2: local convolution + candidate selector (#35371)

Co-authored-by: Jian Chen <jianchen0311@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>

test/registered/unit/model_executor/model_runner_components/test_spec_aux_hidden_state.py

@@ -0,0 +1,38 @@
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+from sglang.srt.model_executor.model_runner_components.spec_aux_hidden_state import (
+    _map_muse_target_layer_ids,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+@pytest.mark.parametrize(
+    ("target_model_type", "draft_architecture", "expected"),
+    [
+        ("muse_glimmer", "MuseGlimmerAssistantModel", [2, 14, 26, 38, 50]),
+        ("muse_glimmer", "DFlash2DraftModel", [2, 14, 26, 38, 50]),
+        ("muse_glimmer", "DFlashDraftModel", [1, 13, 25, 37, 49]),
+        ("qwen3", "DFlash2DraftModel", [1, 13, 25, 37, 49]),
+        ("qwen3", "MuseGlimmerAssistantModel", [1, 13, 25, 37, 49]),
+    ],
+)
+def test_muse_target_layer_id_mapping(target_model_type, draft_architecture, expected):
+    """The +1 belongs to Muse targets, which report layer outputs where the rest
+    report layer inputs. The draft architecture alone does not earn it."""
+    assert (
+        _map_muse_target_layer_ids(
+            target_hf_config=SimpleNamespace(model_type=target_model_type),
+            draft_hf_config=SimpleNamespace(architectures=[draft_architecture]),
+            layer_ids=[1, 13, 25, 37, 49],
+        )
+        == expected
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/spec/test_dflash_logits.py

@@ -0,0 +1,118 @@
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.srt.models.dflash import (
+    CandidateSelector,
+    DFlash2DraftModel,
+    _grouped_conv,
+)
+from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+def test_dflash_unary_logit_transform():
+    logits = torch.tensor([[-100.0, 0.0, 100.0]], dtype=torch.bfloat16)
+    for fields in ({}, {"output_multiplier": 0.2, "final_logit_softcapping": 20.0}):
+        config = parse_dflash_draft_config(
+            draft_hf_config={
+                "num_hidden_layers": 5,
+                "dflash_config": {
+                    "selector_rank": 256,
+                    "selector_top_k": 16,
+                    **fields,
+                },
+            }
+        )
+        actual = DFlash2DraftModel._transform_unary_logits(
+            SimpleNamespace(draft_config=config), logits
+        )
+        expected = logits.float() * config.output_multiplier
+        if config.final_logit_softcapping is not None:
+            expected = torch.tanh(expected / config.final_logit_softcapping)
+            expected *= config.final_logit_softcapping
+        torch.testing.assert_close(actual, expected)
+
+
+def test_selector_greedy_row_walk_is_deterministic_in_a_mixed_batch():
+    """A greedy row walks the argmax, so the q it hands verify has to be the point
+    mass there. Greedy reaches the selector as top_k=1 with the temperature reset
+    to 1.0, so a softmax q stays a real distribution and verify would
+    rejection-sample a deterministic request against it. The row must also not
+    depend on who else is in the batch."""
+    selector = CandidateSelector(hidden_size=4, vocab_size=16, state_rank=2, top_k=4)
+    torch.manual_seed(1)
+    candidate_ids = torch.randint(0, 16, (2, 3, 4))
+    scores = torch.randn(2, 3, 4, 4)
+    uniforms = torch.tensor([[0.2, 0.7, 0.4], [0.8, 0.1, 0.6]])
+    temperatures = torch.tensor([1.0, 0.7])
+    greedy_mask = torch.tensor([True, False])
+
+    mixed_tokens, mixed_q = selector.sample_path(
+        candidate_ids=candidate_ids,
+        scores=scores,
+        uniforms=uniforms,
+        temperatures=temperatures,
+        greedy_mask=greedy_mask,
+    )
+    assert torch.all((mixed_q[0] == 0) | (mixed_q[0] == 1))
+    for row in range(2):
+        tokens, q_rows = selector.sample_path(
+            candidate_ids=candidate_ids[row : row + 1],
+            scores=scores[row : row + 1],
+            uniforms=uniforms[row : row + 1],
+            temperatures=temperatures[row : row + 1],
+            greedy_mask=greedy_mask[row : row + 1],
+        )
+        torch.testing.assert_close(mixed_tokens[row], tokens[0])
+        torch.testing.assert_close(mixed_q[row], q_rows[0])
+
+
+def test_selector_rejects_a_quantized_target_lm_head():
+    """The candidate matmuls read the lm_head weight directly, so a packed or
+    absent weight would be read as if it were dense."""
+    model = SimpleNamespace(
+        lm_head=SimpleNamespace(weight=torch.empty(8, 4, dtype=torch.int8)),
+        candidate_selector=SimpleNamespace(top_k=4),
+    )
+    with pytest.raises(RuntimeError, match="requires a dense"):
+        DFlash2DraftModel.compute_candidates(model, torch.randn(2, 4))
+
+
+def test_grouped_conv_supports_runtime_block_sizes():
+    """The conv indexes a position inside the block, so it must follow whatever
+    block size the worker resolved -- including one that is not a power of two."""
+    torch.manual_seed(0)
+    groups, group_size, taps = 3, 2, 2
+    hidden_size = groups * group_size
+    batch_size = 2
+
+    for block_size in (5, 8, 16):
+        hidden = torch.randn(batch_size * block_size, hidden_size)
+        delta = torch.randn(batch_size * block_size, taps, groups)
+        base = torch.randn(taps, hidden_size)
+
+        actual = _grouped_conv(
+            hidden, delta, base, block_size, groups, group_size, taps
+        )
+
+        expected = torch.empty_like(hidden)
+        hidden_3d = hidden.view(batch_size, block_size, groups, group_size)
+        delta_4d = delta.view(batch_size, block_size, taps, groups)
+        base_3d = base.view(taps, groups, group_size)
+        for batch in range(batch_size):
+            for position in range(block_size):
+                value = torch.zeros(groups, group_size)
+                for tap in range(min(taps, position + 1)):
+                    coefficient = base_3d[tap] + delta_4d[batch, position, tap, :, None]
+                    value += coefficient * hidden_3d[batch, position - tap]
+                expected[batch * block_size + position] = value.flatten()
+        torch.testing.assert_close(actual, expected)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))