feat(dspark): add sampling-aware L1 loss

aspec/config/train_config.py

@@ -251,6 +251,13 @@ class TrainingConfig:
     dspark_loss_decay_gamma: float = 4.0
     dspark_ce_loss_alpha: float = 0.1
     dspark_l1_loss_alpha: float = 0.9
+    # Independent sampling transforms for the DSpark L1 distributions.
+    dspark_l1_target_temperature: float = 1.0
+    dspark_l1_target_top_p: float = 1.0
+    dspark_l1_target_top_k: int = -1
+    dspark_l1_draft_temperature: float = 1.0
+    dspark_l1_draft_top_p: float = 1.0
+    dspark_l1_draft_top_k: int = -1
     dspark_confidence_head_alpha: float = 1.0
 
 

aspec/models/dspark.py

@@ -85,6 +85,8 @@ def __init__(
         deterministic_anchor_sampling: bool = False,
         anchor_seed: int = 0,
         anchor_sampling_scope: str = "row",
+        l1_target_sampling: tuple[float, float, int] = (1.0, 1.0, -1),
+        l1_draft_sampling: tuple[float, float, int] = (1.0, 1.0, -1),
     ):
         # Reuse DFlash anchor/mask/noise machinery. The "decay" objective drives
         # the per-within-block position weighting shared by all DSpark terms.
@@ -104,6 +106,8 @@ def __init__(
         self.ce_loss_alpha = float(ce_loss_alpha)
         self.l1_loss_alpha = float(l1_loss_alpha)
         self.confidence_head_alpha = float(confidence_head_alpha)
+        self.l1_sampling = (l1_target_sampling, l1_draft_sampling)
+        self.l1_sampling_enabled = self.l1_sampling != ((1.0, 1.0, -1), (1.0, 1.0, -1))
         self.l1_loss_backend = None
         self.l1_kernel_invocations = {}
         self._use_l1_kernel_optimizations = True
@@ -550,7 +554,8 @@ def _prediction_terms(self, hidden, target_ids, prev_ids, head_weight,
         accept_rate = None
         if aligned_target_hidden is not None:
             indexed_teacher = (
-                self._use_l1_kernel_optimizations and cached_target_positions is not None
+                not self.l1_sampling_enabled
+                and self._use_l1_kernel_optimizations and cached_target_positions is not None
                 and use_sm90_tiled_l1(logits, cached_target_logits, True)
             )
             target_indices = None
@@ -570,7 +575,9 @@ def _prediction_terms(self, hidden, target_ids, prev_ids, head_weight,
                 target_logits = cached_target_logits if cached_target_logits is not None else F.linear(
                     aligned_target_hidden.reshape(-1, aligned_target_hidden.shape[-1]), head_weight
                 ).view_as(logits)
-            if indexed_teacher or use_sm90_tiled_l1(logits, target_logits):
+            if not self.l1_sampling_enabled and (
+                indexed_teacher or use_sm90_tiled_l1(logits, target_logits)
+            ):
                 from aspec.models.ops.tiled_l1 import tiled_l1_from_logits
 
                 l1 = tiled_l1_from_logits(logits, target_logits, target_indices=target_indices,
@@ -579,7 +586,8 @@ def _prediction_terms(self, hidden, target_ids, prev_ids, head_weight,
                 kernel = ("indexed_" if indexed_teacher else "") + self.l1_loss_backend
                 self.l1_kernel_invocations[kernel] = self.l1_kernel_invocations.get(kernel, 0) + 1
             else:
-                l1 = compiled_l1_from_logits(logits, target_logits)
+                l1 = (compiled_l1_from_logits(logits, target_logits, *self.l1_sampling)
+                      if self.l1_sampling_enabled else compiled_l1_from_logits(logits, target_logits))
                 self.l1_loss_backend = "torch_compiled"
             if self.l1_loss_alpha > 0:
                 l1_num = (l1 * weights).sum()

aspec/models/ops/loss.py

@@ -56,11 +56,38 @@ def _softmax_from_logits(logits: torch.Tensor) -> torch.Tensor:
     return torch.exp(logits_f32 - torch.logsumexp(logits_f32, dim=-1, keepdim=True))
 
 
+def sampling_logits(logits: torch.Tensor, sampling=(1.0, 1.0, -1)) -> torch.Tensor:
+    """Differentiable logits on joint top-k/top-p support; ties use token-ID order.
+
+    Both filters use the full temperature-scaled distribution, matching the
+    pinned SGLang joint-filter convention. The discrete support has no gradient.
+    """
+    if sampling == (1.0, 1.0, -1):
+        return logits
+    temperature, top_p, top_k = sampling
+    scaled = logits.float()
+    if temperature != 1.0:
+        scaled = (scaled - scaled.amax(dim=-1, keepdim=True)) / temperature
+    vocab = scaled.shape[-1]
+    if top_p == 1.0 and (top_k == -1 or top_k >= vocab):
+        return scaled
+    sorted_logits, indices = scaled.sort(dim=-1, descending=True, stable=True)
+    remove = torch.zeros_like(sorted_logits, dtype=torch.bool)
+    if top_p < 1.0:
+        probs = sorted_logits.softmax(dim=-1)
+        remove = probs.cumsum(dim=-1) - probs > top_p
+    if 0 < top_k < vocab:
+        remove = remove | (torch.arange(vocab, device=logits.device) >= top_k)
+    return scaled.masked_fill(remove.scatter(-1, indices, remove), float("-inf"))
+
+
 @torch.compile(dynamic=None)
-def compiled_l1_from_logits(logits: torch.Tensor, target_logits: torch.Tensor) -> torch.Tensor:
-    """Per-token L1 distance between FP32 softmax distributions."""
-    draft_probs = torch.softmax(logits.float(), dim=-1)
-    target_probs = torch.softmax(target_logits.float(), dim=-1)
+def compiled_l1_from_logits(logits: torch.Tensor, target_logits: torch.Tensor,
+                            target_sampling=(1.0, 1.0, -1),
+                            draft_sampling=(1.0, 1.0, -1)) -> torch.Tensor:
+    """Per-token L1 distance between optional sampling-transformed FP32 distributions."""
+    draft_probs = torch.softmax(sampling_logits(logits, draft_sampling).float(), dim=-1)
+    target_probs = torch.softmax(sampling_logits(target_logits, target_sampling).float(), dim=-1)
     return (draft_probs - target_probs).abs().sum(dim=-1)
 
 

aspec/planning.py

@@ -341,6 +341,26 @@ def _load_draft_descriptor(config: DictConfig) -> dict[str, Any]:
     except ValueError as exc:
         raise PlanCompileError(str(exc)) from exc
     backend = _select(config, "training.attention_backend")
+    l1_sampling_enabled = False
+    for side in ("target", "draft"):
+        prefix = f"training.dspark_l1_{side}_"
+        temperature = _select(config, prefix + "temperature", 1.0)
+        top_p = _select(config, prefix + "top_p", 1.0)
+        top_k = _select(config, prefix + "top_k", -1)
+        if (isinstance(temperature, bool) or not isinstance(temperature, (int, float))
+                or not math.isfinite(temperature) or temperature <= 0):
+            raise PlanCompileError(f"{prefix}temperature must be finite and positive")
+        if (isinstance(top_p, bool) or not isinstance(top_p, (int, float))
+                or not math.isfinite(top_p) or not 0 < top_p <= 1):
+            raise PlanCompileError(f"{prefix}top_p must be in (0, 1]")
+        if type(top_k) is not int or (top_k != -1 and top_k < 1):
+            raise PlanCompileError(f"{prefix}top_k must be -1 or a positive integer")
+        l1_sampling_enabled |= (temperature, top_p, top_k) != (1.0, 1.0, -1)
+    if l1_sampling_enabled:
+        indexer_warmup = any((attention.get(kind) or {}).get("training_stage") == "indexer_warmup"
+                            for kind in ("csa", "qsa", "dsa"))
+        if family != "dspark" or _select(config, "training.dspark_l1_loss_alpha") <= 0 or indexer_warmup:
+            raise PlanCompileError("DSpark L1 sampling requires DSpark with positive L1 weight outside indexer warmup")
     eagle3_target_cache = _select(config, "training.eagle3_target_cache", "lazy")
     eagle3_loss_checkpointing = _select(
         config, "training.eagle3_loss_gradient_checkpointing"
@@ -1262,6 +1282,9 @@ def _workload(
                 config, "training.dspark_confidence_head_alpha"
             ),
         }
+        for side in ("target", "draft"):
+            for name, default in (("temperature", 1.0), ("top_p", 1.0), ("top_k", -1)):
+                objective[f"l1_{side}_{name}"] = _select(config, f"training.dspark_l1_{side}_{name}", default)
     else:
         objective = {
             "family": family,

aspec/training/checkpoint.py

@@ -80,6 +80,10 @@ def _logical_training_topology(args: Any, world_size: int) -> dict[str, Any]:
                   "ce_loss_alpha", "l1_loss_alpha")
         if prefix == "dspark":
             fields += ("confidence_head_alpha",)
+            for side in ("target", "draft"):
+                for name, default in (("temperature", 1.0), ("top_p", 1.0), ("top_k", -1)):
+                    key = f"l1_{side}_{name}"
+                    topology["training_objective"][key] = getattr(args, f"dspark_{key}", default)
         else:
             fields += ("loss_objective", "dpace_alpha")
         for name in fields:
@@ -140,6 +144,13 @@ def _validate_exact_resume_topology(args: Any, metadata: dict[str, Any], *, chec
     if saved_topology is not None and not isinstance(saved_topology, dict):
         raise RuntimeError("Exact resume checkpoint has invalid training_topology metadata")
     current_topology = _logical_training_topology(args, current_world_size)
+    if getattr(args, "draft_algorithm", None) == "dspark":
+        saved_objective = (saved_topology or {}).get("training_objective") or {}
+        for side in ("target", "draft"):
+            for name, default in (("temperature", 1.0), ("top_p", 1.0), ("top_k", -1)):
+                key = f"l1_{side}_{name}"
+                if saved_objective.get(key, default) != getattr(args, f"dspark_{key}", default):
+                    raise RuntimeError(f"Exact resume cannot change dspark_{key}; use weights-only initialization")
     # Checkpoints predating this option always sampled per physical row.
     # Missing semantic identity must not authorize a switch to per-sample training.
     if (getattr(args, "draft_algorithm", None) in {"dspark", "dflash", "dflash2"}

aspec/training/dspark_trainer.py

@@ -80,6 +80,16 @@ def _build_training_wrapper(self, draft_model):
             loss_decay_gamma=self.loss_decay_gamma,
             ce_loss_alpha=self.ce_loss_alpha,
             l1_loss_alpha=self.l1_loss_alpha,
+            l1_target_sampling=(
+                getattr(self.args, "dspark_l1_target_temperature", 1.0),
+                getattr(self.args, "dspark_l1_target_top_p", 1.0),
+                getattr(self.args, "dspark_l1_target_top_k", -1),
+            ),
+            l1_draft_sampling=(
+                getattr(self.args, "dspark_l1_draft_temperature", 1.0),
+                getattr(self.args, "dspark_l1_draft_top_p", 1.0),
+                getattr(self.args, "dspark_l1_draft_top_k", -1),
+            ),
             confidence_head_alpha=self.confidence_head_alpha,
             deterministic_anchor_sampling=getattr(
                 self.args, "deterministic_anchor_sampling", False

docs/user-guide/training.md

@@ -179,3 +179,25 @@ mooncake:
 Run `aspec --config offline.yaml --dry-run`, then `aspec --config offline.yaml`. CPU actors replay the features while the Target checkpoint supplies frozen embeddings and the teacher head. Evaluation uses the saved eval split; checkpoint and exact-resume commands are unchanged.
 
 Dataset metadata binds the Target configuration, tokenizer and weights by SHA256, plus the feature layers, dtype, layout and normalization convention. Replay validates sample checksums and tensor shapes. Set `training.max_seq_length` to cover all recorded sequences; changing preprocessing requires regenerating the dataset. The dataset identity participates in exact-resume validation.
+
+## DSpark L1 sampling distributions
+
+DSpark L1 supports independent Target and Draft sampling parameters. Both sides default to `temperature=1.0`, `top_p=1.0` and `top_k=-1` (disabled). Configure either side or both:
+
+```yaml
+training:
+  dspark_l1_target_temperature: 0.8
+  dspark_l1_target_top_p: 0.95
+  dspark_l1_target_top_k: 50
+  dspark_l1_draft_temperature: 1.0
+  dspark_l1_draft_top_p: 0.9
+  dspark_l1_draft_top_k: 40
+```
+
+Each side computes $p=\operatorname{softmax}(z/T)$ in FP32. Top-k and top-p select from this full distribution, then renormalize their intersection, following the pinned SGLang joint-filter convention. Ties use increasing token ID. Top-p retains tokens whose preceding cumulative mass is at most top-p. Temperature must be finite and positive, top-p in $(0,1]$, and top-k `-1` or a positive integer; top-k at least the vocabulary size keeps every token.
+
+The transformed distributions define $L_{\mathrm{L1}}=\sum_i|P_s(i)-Q_s(i)|$ and the confidence label $\alpha=\operatorname{stopgrad}(1-L_{\mathrm{L1}}/2)$. Acceptance metrics use the same distributions. CE uses the original Draft logits, including Markov-head correction. Loss coefficients and masks are unchanged, with no $T^2$ multiplier. These parameters are independent of `decode.*` and apply only to DSpark with positive L1 weight outside indexer warmup.
+
+Draft gradients flow through retained logits and their normalization. Discrete selection has no gradient: top-k one, or disjoint Target/Draft supports, gives zero L1 gradient while the supports stay fixed. CE retains its own gradients.
+
+Default settings retain the existing kernels; non-default settings use the `torch_compiled` L1 path. LM-head computation is unchanged. Exact resume requires the same sampling settings, treating missing fields in older checkpoints as defaults; use weights-only initialization to change them.