ci: add DSpark and DFlash2 GPU smoke coverage

.github/workflows/gpu-integration.yml

@@ -103,7 +103,7 @@ jobs:
             });
 
   integration:
-    name: Qwen3.8 Eagle3 + DSpark (GB200, 2 GPUs)
+    name: Qwen3.8 Eagle3 + DSpark + DFlash2 (GB200, 2 GPUs)
     needs: authorize
     runs-on: ${{ vars.TORCHSPEC_GPU_DISPATCHER_LABEL }}
     timeout-minutes: 720

.github/workflows/gpu-nightly.yml

@@ -58,7 +58,7 @@ jobs:
           } >> "$GITHUB_STEP_SUMMARY"
 
   integration:
-    name: Qwen3.8 full-epoch GPU training
+    name: Qwen3.8 Eagle3 + DSpark + DFlash2 GPU training
     needs: changes
     if: needs.changes.outputs.should_run == 'true'
     runs-on: ${{ vars.TORCHSPEC_GPU_DISPATCHER_LABEL }}

configs/ci/qwen3_8_27b_dflash2_draft.json

@@ -0,0 +1,37 @@
+{
+  "architectures": ["DFlash2DraftModel"],
+  "model_type": "qwen3_dflash2",
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "hidden_size": 5120,
+  "intermediate_size": 8192,
+  "num_hidden_layers": 5,
+  "num_attention_heads": 40,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "vocab_size": 248320,
+  "rms_norm_eps": 1e-6,
+  "max_position_embeddings": 262144,
+  "rope_theta": 10000000.0,
+  "num_target_layers": 5,
+  "target_hidden_size": 5120,
+  "target_num_hidden_layers": 64,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "is_causal": false,
+  "dflash_config": {
+    "block_size": 8,
+    "conv_kernel_size": 2,
+    "conv_group_size": 16,
+    "selector_rank": 256,
+    "selector_top_k": 16,
+    "target_layer_ids": [3, 15, 31, 47, 59],
+    "mask_token_id": 248044
+  },
+  "tie_word_embeddings": false
+}

configs/ci/qwen3_8_27b_dspark_draft.json

@@ -3,7 +3,7 @@
   "model_type": "qwen3_dspark",
   "hidden_size": 5120,
   "intermediate_size": 8192,
-  "num_hidden_layers": 1,
+  "num_hidden_layers": 5,
   "num_attention_heads": 40,
   "num_key_value_heads": 8,
   "head_dim": 128,
@@ -16,7 +16,7 @@
   "target_num_hidden_layers": 64,
   "target_layer_ids": [3, 15, 31, 47, 59],
   "mask_token_id": 248044,
-  "markov_rank": 64,
+  "markov_rank": 256,
   "markov_head_type": "vanilla",
   "enable_confidence_head": true,
   "confidence_head_with_markov": true,

configs/ci/vllm_qwen3_8_27b_dflash2_2gpu_smoke.yaml

@@ -0,0 +1,81 @@
+# Two-GPU CI lane: the same Qwen3.8-27B target on one GB200 and a
+# production-depth DFlash2 draft on the other. This exercises block diffusion,
+# grouped convolutions, candidate selection, and selector-loss training.
+model:
+  target_model_path: Qwen/Qwen3.8-27B
+  target_model_backend: vllm
+  draft_model_config: configs/ci/qwen3_8_27b_dflash2_draft.json
+  embedding_key: model.language_model.embed_tokens.weight
+  lm_head_key: lm_head.weight
+  norm_key: model.language_model.norm.weight
+  trust_remote_code: true
+
+dataset:
+  train_data_path: ../../examples/data/sample_conversations.jsonl
+  chat_template: qwen
+  prompt_key: conversations
+  shuffle_dataset: false
+  length_group_size: 1
+  min_loss_tokens: 32
+
+training:
+  attention_backend: flex_attention
+  micro_batch_size: 1
+  draft_accumulation_steps: 1
+  learning_rate: 1e-4
+  max_concurrent_batches: 1
+  max_grad_norm: 0.5
+  # Keep the complete 12,913-token record so DFlash2 is exercised after
+  # multiple vLLM chunked-prefill iterations.
+  max_seq_length: 16384
+  num_epochs: 1
+  prefetch_depth: 1
+  save_interval: 0
+  save_per_epoch: false
+  seed: 42
+  training_num_gpus_per_node: 1
+  training_num_nodes: 1
+  warmup_ratio: 0.0
+  dflash_block_size: 8
+  dflash_num_anchors: 512
+  dflash_num_target_layers: 5
+  dflash_loss_decay_gamma: 7.0
+  dflash2_selector_loss_alpha: 1.0
+
+inference:
+  inference_engine_type: vllm
+  store_last_hidden_states: false
+  inference_num_gpus: 1
+  inference_num_gpus_per_engine: 1
+  inference_num_gpus_per_node: 2
+  inference_batch_size: 1
+  inference_buffer_threshold: 2
+  max_sample_pool_size: 4
+  # These must match dflash_config.target_layer_ids in the draft config.
+  aux_hidden_states_layers: [3, 15, 31, 47, 59]
+  vllm:
+    tp_size: 1
+    mem_fraction_static: 0.55
+    extra_args:
+      enforce_eager: true
+      language_model_only: true
+      max_model_len: 16384
+      max_num_batched_tokens: 4096
+
+mooncake:
+  master_server_address: null
+  metadata_server: null
+  protocol: tcp
+  global_segment_size: 4GB
+  local_buffer_size: 1GB
+
+output_dir: ""
+cache_dir: ../../.ci-cache/qwen3-8-27b-dflash2-smoke
+model_download_dir: null
+
+logging:
+  report_to: none
+
+debug:
+  enable_perf_metrics: true
+  save_debug_train_data: null

configs/ci/vllm_qwen3_8_27b_dspark_2gpu_smoke.yaml

@@ -1,6 +1,6 @@
 # Two-GPU CI lane: the same Qwen3.8-27B target on one GB200 and a
-# one-layer DSpark draft on the other. This exercises block diffusion,
-# Markov, confidence, CE, and L1 training paths without a production-sized draft.
+# production-depth DSpark draft on the other. This exercises block diffusion,
+# Markov, confidence, CE, and L1 training paths with the standard DSpark shape.
 model:
   target_model_path: Qwen/Qwen3.8-27B
   target_model_backend: vllm
@@ -16,7 +16,7 @@ dataset:
   prompt_key: conversations
   shuffle_dataset: false
   length_group_size: 1
-  min_loss_tokens: 4
+  min_loss_tokens: 32
 
 training:
   attention_backend: flex_attention
@@ -36,8 +36,8 @@ training:
   training_num_gpus_per_node: 1
   training_num_nodes: 1
   warmup_ratio: 0.0
-  dflash_block_size: 2
-  dspark_num_anchors: 4
+  dflash_block_size: 7
+  dspark_num_anchors: 512
   dspark_num_target_layers: 5
   dspark_loss_decay_gamma: 4.0
   dspark_ce_loss_alpha: 0.1
@@ -80,5 +80,5 @@ logging:
   report_to: none
 
 debug:
-  enable_perf_metrics: false
+  enable_perf_metrics: true
   save_debug_train_data: null

configs/ci/vllm_qwen3_8_27b_eagle3_2gpu_smoke.yaml

@@ -74,5 +74,5 @@ logging:
   report_to: none
 
 debug:
-  enable_perf_metrics: false
+  enable_perf_metrics: true
   save_debug_train_data: null

tests/test_dflash2.py

@@ -35,6 +35,7 @@
     _convert_fsdp_to_hf,
     _fixup_export_config,
 )
+from torchspec.config import load_config
 from torchspec.models.dflash import _create_dflash_mask_mod
 from torchspec.models.dflash2 import DFlash2Model
 from torchspec.models.draft.auto import AutoDraftModelConfig
@@ -171,6 +172,29 @@ def test_repository_config_dispatches_to_dflash2(self):
         with mock.patch.dict("sys.modules", {"torchspec.training.dflash2_trainer": trainer_module}):
             self.assertIs(_trainer_class_for_config(config), trainer_class)
 
+    def test_ci_config_dispatches_to_dflash2_with_matching_capture_layers(self):
+        config_path = ROOT / "configs" / "ci" / "vllm_qwen3_8_27b_dflash2_2gpu_smoke.yaml"
+        config = load_config(str(config_path))
+        draft_config = AutoDraftModelConfig.from_file(config.model.draft_model_config)
+
+        self.assertIsInstance(draft_config, DFlash2Config)
+        self.assertEqual(draft_config.architectures, ["DFlash2DraftModel"])
+        self.assertEqual(
+            list(config.inference.aux_hidden_states_layers),
+            draft_config.target_layer_ids,
+        )
+        self.assertEqual(config.training.dflash_block_size, draft_config.block_size)
+        self.assertEqual(config.training.dflash_num_anchors, 512)
+        self.assertEqual(config.dataset.min_loss_tokens, 32)
+        self.assertEqual(
+            config.training.dflash_num_target_layers,
+            draft_config.num_target_layers,
+        )
+        self.assertEqual(draft_config.num_hidden_layers, 5)
+        self.assertEqual(draft_config.target_num_hidden_layers, 64)
+        self.assertEqual(draft_config.selector_rank, 256)
+        self.assertEqual(draft_config.selector_top_k, 16)
+
     def test_legacy_dflash_dispatch_is_unchanged(self):
         config = AutoDraftModelConfig.from_dict(
             {

tests/test_dspark.py

@@ -37,9 +37,11 @@
 
 import math
 import unittest
+from pathlib import Path
 
 import torch
 
+from torchspec.config import load_config
 from torchspec.models.draft.auto import AutoDraftModelConfig, AutoEagle3DraftModel
 from torchspec.models.draft.dflash import DFlashConfig
 from torchspec.models.draft.dspark import (
@@ -54,6 +56,7 @@
 from torchspec.models.dspark import DSparkModel
 
 CE_A, L1_A, CF_A = 0.1, 0.9, 1.0
+ROOT = Path(__file__).resolve().parents[1]
 
 
 def _make_dspark_config(
@@ -135,6 +138,23 @@ def test_subclasses_dflash_and_attrs(self):
         self.assertTrue(cfg.enable_confidence_head)
         self.assertFalse(cfg.fc_norm)
 
+    def test_ci_config_uses_standard_dspark_shape(self):
+        config_path = ROOT / "configs" / "ci" / "vllm_qwen3_8_27b_dspark_2gpu_smoke.yaml"
+        config = load_config(str(config_path))
+        draft_config = AutoDraftModelConfig.from_file(config.model.draft_model_config)
+
+        self.assertIsInstance(draft_config, DSparkConfig)
+        self.assertEqual(config.training.dflash_block_size, 7)
+        self.assertEqual(config.training.dspark_num_anchors, 512)
+        self.assertEqual(config.dataset.min_loss_tokens, 32)
+        self.assertEqual(draft_config.num_hidden_layers, 5)
+        self.assertEqual(draft_config.markov_rank, 256)
+        self.assertEqual(draft_config.target_num_hidden_layers, 64)
+        self.assertEqual(
+            list(config.inference.aux_hidden_states_layers),
+            draft_config.target_layer_ids,
+        )
+
     def test_optional_fc_norm_normalizes_each_target_layer_before_projection(self):
         cfg = _make_dspark_config(H=16, num_target_layers=3, fc_norm=True)
         model = DSparkDraftModel(cfg).to(dtype=torch.float32)

tools/ci/gpu_2gpu.sbatch

@@ -24,6 +24,7 @@ srun --gpus=2 \
     export RUNNER_TEMP=/ci-tmp
     export HF_HOME=/ci-tmp/huggingface
     export PYTHONPATH="/workspace${PYTHONPATH:+:${PYTHONPATH}}"
+    export VLLM_USE_V2_MODEL_RUNNER=1
     export TORCHSPEC_CI_ARTIFACT_DIR=/artifacts
     export TORCHSPEC_CI_MODEL_PATH="/model-cache/snapshots/${TORCHSPEC_CI_MODEL_REVISION}"
     export TORCHSPEC_CI_MODEL_CACHE=/ci-tmp/huggingface