Upload folder using huggingface_hub

luau_infer.py

@@ -626,6 +626,19 @@ class Engine:
         return st["res"]
 
 
+class Unfusable(Exception):
+    """Settings the fused decoder can't reproduce, found after the prompt pass."""
+
+
+def _observes_only(processors, input_ids: Tensor, logits: Tensor, first: Tensor) -> bool:
+    """True if ``processors`` leave the scores unchanged, at the prompt and with generated tokens."""
+    fake = torch.cat((input_ids, first[:, None].expand(-1, 8)), 1)
+    for ids in (input_ids, fake):
+        if not torch.equal(processors(ids, logits.clone()), logits):
+            return False
+    return True
+
+
 class _Lookup:
     """Host mirror of the engine's draft for one row: estimates whether speculation pays."""
 
@@ -701,7 +714,8 @@ class FusedDecoder:
     def decode(self, input_ids: Tensor, attention_mask: Tensor | None, *, max_new_tokens: int,
                sample: bool = False, temperature: float = 1.0, top_k: int = 0, top_p: float = 1.0,
                min_p: float = 0.0, repetition_penalty: float = 1.0, num_samples: int = 1, eos=None,
-               pad: int | None = None, stop=(), tokenizer=None, streamer=None, stopping_criteria=None) -> Tensor:
+               pad: int | None = None, stop=(), tokenizer=None, streamer=None, stopping_criteria=None,
+               logits_processor=None) -> Tensor:
         """Prompt + completion ids ``(B * num_samples, T0 + n)``, ``pad`` after a row's end."""
         eng = self.engine
         if attention_mask is None:
@@ -725,6 +739,8 @@ class FusedDecoder:
         eng.allocate(B, max(T0 + n_new + (self.DEPTH + 3) * (K + 1), self.MIN_CACHE if B <= 8 else 0))
         eng.set_eos(eos_set)
         first = eng.prefill(input_ids, attention_mask, sp)
+        if logits_processor and not _observes_only(logits_processor, input_ids, eng.last_logits, first):
+            raise Unfusable("logits processors that change the scores")
         crit_ids = None
         if stopping_criteria:
             crit_ids = torch.full((B, T0 + n_new + K + 1), pad, dtype=input_ids.dtype, device=input_ids.device)

modeling_luau_coder.py

@@ -1426,8 +1426,6 @@ def _fused_fallback_reason(model, inputs, gc, model_kwargs, logits_processor, st
     used = [k for k, v in extra.items() if v is not None and v is not False]
     if used:
         return f"unsupported generate() arguments {used}"
-    if logits_processor:
-        return "custom logits processors"
     if set(model_kwargs) - _FUSED_MODEL_KWARGS:
         return f"model kwargs {sorted(set(model_kwargs) - _FUSED_MODEL_KWARGS)}"
     for name, ok in _FUSED_NEUTRAL.items():
@@ -1496,15 +1494,20 @@ class _FusedGenerationMixin:
         n_new = gc.max_new_tokens if gc.max_new_tokens is not None else (gc.max_length or 20) - inputs.shape[1]
         stop = gc.stop_strings or []
         stop = [stop] if isinstance(stop, str) else list(stop)
+        from .luau_infer import Unfusable
         decoder = self._fused_decoder(spec)
-        with torch.inference_mode():
-            return decoder.decode(
-                inputs.to(self.device), mask.to(self.device), max_new_tokens=max(0, int(n_new)),
-                sample=bool(gc.do_sample), temperature=gc.temperature or 1.0, top_k=gc.top_k or 0,
-                top_p=1.0 if gc.top_p is None else gc.top_p, min_p=gc.min_p or 0.0,
-                repetition_penalty=gc.repetition_penalty or 1.0, num_samples=gc.num_return_sequences or 1,
-                eos=eos, pad=pad, stop=stop, tokenizer=model_kwargs.get("tokenizer"), streamer=streamer,
-                stopping_criteria=stopping_criteria or None)
+        try:
+            with torch.inference_mode():
+                return decoder.decode(
+                    inputs.to(self.device), mask.to(self.device), max_new_tokens=max(0, int(n_new)),
+                    sample=bool(gc.do_sample), temperature=gc.temperature or 1.0, top_k=gc.top_k or 0,
+                    top_p=1.0 if gc.top_p is None else gc.top_p, min_p=gc.min_p or 0.0,
+                    repetition_penalty=gc.repetition_penalty or 1.0, num_samples=gc.num_return_sequences or 1,
+                    eos=eos, pad=pad, stop=stop, tokenizer=model_kwargs.get("tokenizer"), streamer=streamer,
+                    stopping_criteria=stopping_criteria or None, logits_processor=logits_processor or None)
+        except Unfusable as e:
+            self._fused_last_reason = str(e)
+            return None
 
     def _fused_decoder(self, spec: int):
         """The decoder of these weights, kept between calls (rebuilt when a parameter changed)."""