Add DSpark verification capacity modes

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Author: LucasWilkinson

tests/engine/test_arg_utils.py

@@ -225,6 +225,43 @@ def test_jit_monitor_mode_arg(mode):
     assert engine_args.create_observability_config().jit_monitor_mode == mode
 
 
+def test_dspark_capacity_verification_mode_arg():
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    args = parser.parse_args(
+        [
+            "--spec-method",
+            "ngram",
+            "--spec-tokens",
+            "1",
+            "--dspark-capacity-verification-mode",
+            "mask",
+        ]
+    )
+
+    engine_args = EngineArgs.from_cli_args(args)
+    assert engine_args.dspark_capacity_verification_mode == "mask"
+    speculative_config = engine_args.create_speculative_config(None, None)
+    assert speculative_config is not None
+    assert speculative_config.dspark_capacity_verification_mode == "mask"
+
+
+def test_dspark_capacity_verification_mode_conflicts_with_speculative_config():
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    args = parser.parse_args(
+        [
+            "--speculative-config",
+            '{"method":"ngram","num_speculative_tokens":1,'
+            '"dspark_capacity_verification_mode":"mask"}',
+            "--dspark-capacity-verification-mode",
+            "varlen",
+        ]
+    )
+
+    engine_args = EngineArgs.from_cli_args(args)
+    with pytest.raises(ValueError, match="dspark_capacity_verification_mode"):
+        engine_args.create_speculative_config(None, None)
+
+
 def test_hf_token_get_kwargs():
     kwargs = get_kwargs(ModelConfig)["hf_token"]
 

tests/test_config.py

@@ -1572,6 +1572,55 @@ def test_draft_sample_method_gumbel_is_rejected():
         )
 
 
+def test_dspark_capacity_config_validation():
+    speculative_config = SpeculativeConfig(
+        method="ngram",
+        num_speculative_tokens=1,
+        dspark_confidence_threshold=0.25,
+        dspark_budget_frac=0.5,
+        dspark_capacity_verification_mode="mask",
+    )
+    assert speculative_config.dspark_confidence_threshold == 0.25
+    assert speculative_config.dspark_budget_frac == 0.5
+    assert speculative_config.dspark_capacity_verification_mode == "mask"
+    assert (
+        SpeculativeConfig(
+            method="ngram", num_speculative_tokens=1
+        ).dspark_capacity_verification_mode
+        == "varlen"
+    )
+    assert (
+        SpeculativeConfig(
+            method="ngram",
+            num_speculative_tokens=1,
+            dspark_capacity_verification_mode="compact",
+        ).dspark_capacity_verification_mode
+        == "varlen"
+    )
+    assert (
+        SpeculativeConfig(
+            method="ngram", num_speculative_tokens=1
+        ).dspark_confidence_threshold
+        == 0.0
+    )
+
+    for threshold in (-0.1, 1.1):
+        with pytest.raises(ValueError, match="dspark_confidence_threshold"):
+            SpeculativeConfig(
+                method="ngram",
+                num_speculative_tokens=1,
+                dspark_confidence_threshold=threshold,
+            )
+
+    for budget_frac in (0.0, -0.1, 1.1):
+        with pytest.raises(ValueError, match="dspark_budget_frac"):
+            SpeculativeConfig(
+                method="ngram",
+                num_speculative_tokens=1,
+                dspark_budget_frac=budget_frac,
+            )
+
+
 def test_ir_op_priority_default():
     """Test that IR op priority defaults are set correctly."""
     from vllm.config.kernel import IrOpPriorityConfig

tests/v1/spec_decode/test_rejection_sampler_utils.py

@@ -150,7 +150,7 @@ def test_stochastic_rejection_sample(num_speculative_steps: int, temperature: fl
 
     torch.manual_seed(42)
     device = "cuda"
-    num_trials = 10 * VOCAB_SIZE
+    num_trials = 64
 
     target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
     draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
@@ -288,7 +288,7 @@ def test_placeholder_draft_token_rejected():
     """
     torch.manual_seed(0)
     device = "cuda"
-    num_trials = 64
+    num_trials = 10 * VOCAB_SIZE
     K = 1
     temperature = 0.6
 

tests/v1/worker/test_gpu_block_table.py

@@ -5,6 +5,7 @@
 import torch
 
 from vllm.platforms import current_platform
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
 from vllm.v1.worker.gpu.block_table import BlockTables
 
 pytestmark = pytest.mark.skipif(
@@ -130,3 +131,49 @@ def test_block_tables_apply_staged_writes_single_group():
         block_tables.block_tables[0].gpu[0, :2],
         torch.tensor([1, 2], dtype=torch.int32, device=device),
     )
+
+
+def test_compute_slot_mappings_applies_padding_mask():
+    device = torch.device("cuda")
+    block_tables = BlockTables(
+        block_sizes=[16],
+        max_num_reqs=2,
+        max_num_batched_tokens=8,
+        max_num_blocks_per_group=[4],
+        device=device,
+        kernel_block_sizes=[16],
+    )
+
+    block_tables.append_block_ids(
+        req_index=0,
+        new_block_ids=([2],),
+        overwrite=True,
+    )
+    block_tables.append_block_ids(
+        req_index=1,
+        new_block_ids=([3],),
+        overwrite=True,
+    )
+    block_tables.apply_staged_writes()
+
+    idx_mapping = torch.tensor([0, 1], dtype=torch.int32, device=device)
+    query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
+    positions = torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64, device=device)
+    is_padding = torch.tensor(
+        [False, True, False, False, True, False, False, False],
+        dtype=torch.bool,
+        device=device,
+    )
+
+    slot_mappings = block_tables.compute_slot_mappings(
+        idx_mapping,
+        query_start_loc,
+        positions,
+        num_tokens_padded=8,
+        is_padding=is_padding,
+    )
+    torch.accelerator.synchronize()
+
+    assert slot_mappings.cpu().tolist() == [
+        [32, PAD_SLOT_ID, 34, 48, PAD_SLOT_ID, PAD_SLOT_ID, PAD_SLOT_ID, PAD_SLOT_ID]
+    ]