support dflash2

configs/qwen3.6-27b-dflash2.json

@@ -0,0 +1,57 @@
+{
+  "architectures": [
+    "DFlash2DraftModel"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": null,
+  "is_causal": false,
+  "dflash_config": {
+    "block_size": 8,
+    "conv_group_size": 16,
+    "conv_kernel_size": 2,
+    "mask_token_id": 248070,
+    "selector_rank": 256,
+    "selector_top_k": 16,
+    "target_layer_ids": [
+      1,
+      16,
+      31,
+      46,
+      61
+    ]
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 248044,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "initializer_range": 0.02,
+  "intermediate_size": 17408,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention"
+  ],
+  "max_position_embeddings": 262144,
+  "max_window_layers": 5,
+  "model_type": "qwen3",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 5,
+  "num_key_value_heads": 8,
+  "num_target_layers": 64,
+  "pad_token_id": 248044,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 10000000,
+    "rope_type": "default"
+  },
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.8.1",
+  "use_cache": true,
+  "use_sliding_window": true,
+  "vocab_size": 248320
+}

docs/basic_usage/training.md

@@ -176,6 +176,49 @@ mixed layout must be edited explicitly in the draft JSON.
 
 The `eager`, `sdpa`, and `flex_attention` backends support both layouts.
 
+### DFlash 2
+
+DFlash 2 is a draft-architecture variant of DFlash, not a separate capture
+strategy. Keep `training.strategy: dflash` and select it with a draft config
+whose architecture is `DFlash2DraftModel`. The target server still captures
+the same selected hidden states with `--spec-capture-method dflash`.
+
+The DFlash 2 config additionally defines `conv_kernel_size` and
+`conv_group_size` for the local convolution, plus `selector_rank` and
+`selector_top_k` for candidate-path selection. The base draft head uses the
+configured CE/LK/TV objective. The selector always uses categorical CE over the
+target head's strict unary top-k, exactly as it will be used during inference.
+If the gold token is outside that candidate set, the token contributes no
+selector loss; `selector_coverage` reports how often the gold token is present.
+Both objectives receive the configured fixed-decay or D-PACE position weight,
+and their combined numerator is normalized by the sum of valid effective token
+weights rather than by batch or anchor count.
+
+Set `training.dflash2_selector_loss_alpha` to scale the selector objective.
+`training.dflash2_selector_warmup_ratio` keeps that scale at zero for an initial
+fraction of optimizer steps, and `training.dflash2_selector_ramp_ratio` then
+ramps it linearly to the configured value. Both schedule ratios default to
+zero. A newly initialized selector starts as a unary no-op, so enabling DFlash
+2 does not perturb the initial DFlash proposal scores.
+
+The exported computation and parameter names match the public SGLang DFlash 2
+contract, including optional `output_multiplier` and
+`final_logit_softcapping` transforms from `dflash_config`.
+
+The checked-in Qwen3.6-27B recipe owns the full two-GPU local stack: GPU 0 runs
+the target capture server and GPU 1 runs the trainer.
+
+```bash
+specforge train \
+  -c examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash2-disaggregated.yaml \
+  model.target_model_path=/path/to/Qwen3.6-27B
+```
+
+Export the result with `specforge export --to hf`. Serving requires an SGLang
+version that includes DFlash 2 support (SGLang PR #35371); the serving algorithm
+name remains `DFLASH`, and the exported `DFlash2DraftModel` config enables the
+new path automatically.
+
 Domino and DSpark need their projector/head metadata, so they require an
 explicit draft config (or a pretrained warm-start source that contains
 `config.json`). The old Domino parser exposed an optional config flag, but its
@@ -480,7 +523,7 @@ specforge export --to sglang \
 ```
 
 `--to sglang` currently implements the EAGLE3 serving-key contract. Use
-`--to hf` for DFlash, Domino, DSpark, and P-EAGLE model directories. For an
+`--to hf` for DFlash, DFlash 2, Domino, DSpark, and P-EAGLE model directories. For an
 EAGLE-family self-contained Hugging Face directory, provide the target model as
 the source of the frozen embedding when it is absent from the runtime
 checkpoint:

examples/configs/README.md

@@ -279,8 +279,8 @@ Strategy-specific fields should be written only when tuning that objective:
 
 | Strategy | Fields and defaults |
 | --- | --- |
-| EAGLE3 | `training.ttt_length` (`7`), `training.lk_loss_type` (`null`; `lambda` or `alpha`), `training.kl_scale` (`1.0`), `training.kl_decay` (`1.0`) |
-| DFlash / Domino / D-PACE | `training.num_anchors` (`512`), `training.loss_decay_gamma` (`null`), `training.objective_chunk_blocks` (`128`; `0` materializes all objective logits), `training.loss_type` (`dflash`), `training.dpace_alpha` (`0.5`), `training.lambda_base_start` (`1.0`), `training.lambda_base_decay_ratio` (`0.5`) |
+| EAGLE3 | `training.ttt_length` (`7`), `training.lk_loss_type` (`null`; `lambda`, `alpha`, or `tv`), `training.kl_scale` (`1.0`), `training.kl_decay` (`1.0`) |
+| DFlash / DFlash 2 / Domino / D-PACE | `training.num_anchors` (`512`), `training.loss_decay_gamma` (`null`), `training.objective_chunk_blocks` (`128`; `0` materializes all objective logits), `training.loss_type` (`dflash`; fixed decay, or `dpace`; dynamic weighting), DFlash/DFlash 2's `training.lk_loss_type` (`null`; CE, `lambda`, `alpha`, or `tv`), `training.kl_scale` (`1.0`), `training.kl_decay` (`1.0`), DFlash 2's CE selector objective controls `training.dflash2_selector_loss_alpha` (`1.0`), `training.dflash2_selector_warmup_ratio` (`0.0`), and `training.dflash2_selector_ramp_ratio` (`0.0`), `training.dpace_alpha` (`0.5`), `training.lambda_base_start` (`1.0`), `training.lambda_base_decay_ratio` (`0.5`) |
 | DSpark | Token-pooled objective with valid-first-target anchors and distributed ratio telemetry. Configure the shared `training.num_anchors` (`512`), `training.loss_decay_gamma` (`null`; production recipes use `4.0`), and `training.objective_chunk_blocks` (`128`; `0` materializes all objective logits), plus `training.dspark_ce_loss_alpha` (`0.1`), `training.dspark_l1_loss_alpha` (`0.9`), and `training.dspark_confidence_head_alpha` (`1.0`). |
 | P-EAGLE | `training.num_depths` (`8`), `training.down_sample_ratio` (`0.8`), `training.down_sample_ratio_min` (`0.2`), `training.norm_before_residual` (`null`) |
 
@@ -492,6 +492,23 @@ For deeper lifecycle and recovery semantics, see the
 shared DFlash strategy with `training.loss_type: dpace`; D-PACE is an objective
 selection inside the unified trainer, not another training entry.
 
+DFlash 2 treats the token objective and position weighting as independent A/B
+axes. `training.lk_loss_type: null` keeps hard-target CE, `tv` minimizes the
+one-hot total-variation objective, and `lambda` adaptively mixes CE with TV;
+`alpha` is equivalent to CE for these hard targets. Independently,
+`training.loss_type: dflash` uses `loss_decay_gamma`, while `dpace` applies one
+detached dynamic position weight to both the unary and candidate-selector
+objectives. For example, LK-lambda with D-PACE uses:
+
+```yaml
+training:
+  lk_loss_type: lambda
+  kl_scale: 1.0
+  kl_decay: 1.0
+  loss_type: dpace
+  dpace_alpha: 0.5
+```
+
 Evaluation is currently offline-only and pairs `training.eval_interval` with
 `data.eval_hidden_states_path`. Best checkpoints are
 linked as `<run_id>-best`. Offline text EAGLE3 may enable

examples/configs/online/disaggregated/managed-local/qwen3.6-27b-dflash2-disaggregated.yaml

@@ -0,0 +1,82 @@
+model:
+  target_model_path: /disk3/wjp/pretrained_models/Qwen3.6-27B
+  draft_model_config: configs/qwen3.6-27b-dflash2.json
+  target_backend: sglang
+  trust_remote_code: true
+  embedding_key: model.language_model.embed_tokens.weight
+  torch_dtype: bfloat16
+  mask_token_id: 248070
+  sglang_context_length: 8250
+
+data:
+  train_data_path: /disk3/wjp/datasets/perfectblend/qwen3.6-27b/data/sharegpt/full-1.jsonl
+  max_length: 8192
+  chat_template: qwen3.5
+  cache_dir: cache/qwen3.6-27b-pbe-1
+  build_dataset_num_proc: 64
+
+training:
+  # DFlash2 shares DFlash capture/data semantics; the draft architecture turns
+  # on the local convolutions and candidate-selector objective.
+  strategy: dflash
+  num_epochs: 60000
+  max_steps: 100000
+  batch_size: 1
+  accumulation_steps: 4
+  learning_rate: 0.0006
+  warmup_ratio: 0.04
+  max_grad_norm: 1.0
+  attention_backend: flex_attention
+  num_anchors: 512
+  loss_decay_gamma: 7.0
+  lk_loss_type: lambda
+  loss_type: dpace
+  objective_chunk_blocks: 128
+  dflash2_selector_loss_alpha: 1.0
+  dflash2_selector_warmup_ratio: 0.0005
+  dflash2_selector_ramp_ratio: 0.0005
+  save_interval: 200
+  log_interval: 1
+  dist_timeout: 30
+  seed: 42
+
+# tracking:
+#   report_to: wandb
+#   wandb_project: qwen36-dflash2
+#   wandb_name: qwen36-27b-dflash2-1srv-1trainer
+
+tracking:
+  report_to: tensorboard
+
+runtime:
+  producer_lease: 4
+  producer_concurrency: 2
+  in_flight_high_watermark: 100
+  in_flight_low_watermark: 40
+
+run_id: qwen3.6-27b-dflash2-disaggregated-with-selector
+output_dir: outputs/qwen3.6-27b-dflash2-disaggregated-with-selector
+
+deployment:
+  mode: disaggregated
+  trainer:
+    nnodes: 1
+    nproc_per_node: 1
+  disaggregated:
+    control_dir: outputs/qwen3.6-27b-dflash2-disaggregated-with-selector/control
+    consumer_state_dir: outputs/qwen3.6-27b-dflash2-disaggregated-with-selector/consumer-state
+    backend: mooncake
+    managed_local:
+      trainer_cuda_visible_devices: ["5"]
+      mooncake:
+        protocol: tcp
+        global_segment_size_bytes: 68719476736
+        local_buffer_size_bytes: 1073741824
+        default_kv_lease_ttl_ms: 5000
+      capture_servers:
+        - port: 30000
+          cuda_visible_devices:
+            - "2"
+          tp_size: 1
+          mem_fraction_static: 0.8
+

scripts/gates/normalize_dflash_export.py

@@ -14,6 +14,28 @@
     "enable_confidence_head",
     "confidence_head_with_markov",
 )
+_DFLASH2_ARCHITECTURE = "DFlash2DraftModel"
+_DFLASH2_FIELDS = (
+    "conv_group_size",
+    "conv_kernel_size",
+    "selector_rank",
+    "selector_top_k",
+)
+
+
+def _positive_integer(value: Any) -> bool:
+    return isinstance(value, int) and not isinstance(value, bool) and value > 0
+
+
+def _normalize_dflash2(config: Dict[str, Any], method_config: Dict[str, Any]) -> None:
+    for key in _DFLASH2_FIELDS:
+        value = method_config.get(key)
+        if not _positive_integer(value):
+            raise ValueError(
+                f"DFlash2 export requires a positive integer dflash_config.{key}, "
+                f"got {value!r}"
+            )
+    config["architectures"] = [_DFLASH2_ARCHITECTURE]
 
 
 def _normalize_dspark(config: Dict[str, Any], method_config: Dict[str, Any]) -> None:
@@ -69,12 +91,25 @@ def normalize_export(config_path: str, expected_block_size: int) -> Dict[str, An
     with path.open(encoding="utf-8") as handle:
         config = json.load(handle)
 
-    block_size = config.get("block_size")
+    method_config = config.get("dflash_config") or {}
+    top_level_block_size = config.get("block_size")
+    nested_block_size = method_config.get("block_size")
+    if (
+        top_level_block_size is not None
+        and nested_block_size is not None
+        and top_level_block_size != nested_block_size
+    ):
+        raise ValueError(
+            "exported block_size conflict: top-level "
+            f"{top_level_block_size!r} != dflash_config {nested_block_size!r}"
+        )
+    block_size = (
+        top_level_block_size if top_level_block_size is not None else nested_block_size
+    )
     if block_size != expected_block_size:
         raise ValueError(
             f"exported block_size={block_size!r}, expected {expected_block_size}"
         )
-    method_config = config.get("dflash_config") or {}
     projector_type = method_config.get("projector_type", "dflash")
     if projector_type not in {"dflash", "domino", "dspark"}:
         raise ValueError(
@@ -84,6 +119,8 @@ def normalize_export(config_path: str, expected_block_size: int) -> Dict[str, An
 
     if projector_type == "dspark":
         _normalize_dspark(config, method_config)
+    elif _DFLASH2_ARCHITECTURE in (config.get("architectures") or []):
+        _normalize_dflash2(config, method_config)
     else:
         config["architectures"] = ["DFlashDraftModel"]
     config.pop("auto_map", None)