Phase A serving speedup: zero-copy state-cache expand (no deepcopy per decision)

inference/test_parity.py

@@ -41,6 +41,49 @@ NOUL_ITEMS = [
 ]
 
 
+def test_expand_state_cache_matches_deepcopy_and_does_not_alias():
+    """Phase A (serving speedup): `_expand_state_cache` replaces `copy.deepcopy(state_cache)`
+    + `batch_repeat_interleave` with a zero-copy `.expand()`. No checkpoint needed -- this
+    exercises the cache mechanics directly against transformers' real DynamicCache/DynamicLayer
+    with random tensors, checking (a) bit-identical results to the old deepcopy+repeat
+    reference and (b) that expanding+updating the new cache never mutates the original."""
+    import torch
+    from transformers.cache_utils import DynamicCache
+
+    from typical.native import _expand_state_cache
+
+    torch.manual_seed(0)
+    B, H, Ls, D, m = 1, 4, 7, 8, 5
+    k0, v0 = torch.randn(B, H, Ls, D), torch.randn(B, H, Ls, D)
+    k1, v1 = torch.randn(B, H, Ls, D), torch.randn(B, H, Ls, D)
+    cache = DynamicCache(ddp_cache_data=[(k0, v0), (k1, v1)])
+    orig_keys = [layer.keys.clone() for layer in cache.layers]
+    orig_values = [layer.values.clone() for layer in cache.layers]
+
+    # ground truth: the old path (deepcopy, then the real repeat_interleave every layer used
+    # to get regardless of whether it holds K/V or Gated-DeltaNet conv/recurrent state)
+    import copy
+    ref = copy.deepcopy(cache)
+    ref.batch_repeat_interleave(m)
+
+    new = _expand_state_cache(cache, m)
+
+    q_len = 3
+    torch.manual_seed(1)
+    new_k = [torch.randn(m, H, q_len, D), torch.randn(m, H, q_len, D)]
+    new_v = [torch.randn(m, H, q_len, D), torch.randn(m, H, q_len, D)]
+    for i in range(2):
+        k_ref, v_ref = ref.update(new_k[i], new_v[i], i)
+        k_new, v_new = new.update(new_k[i], new_v[i], i)
+        assert torch.equal(k_ref, k_new) and torch.equal(v_ref, v_new), f"layer {i} mismatch"
+
+    # the original (batch=1) cache must be untouched by both the expand and the update above
+    for i, layer in enumerate(cache.layers):
+        assert torch.equal(layer.keys, orig_keys[i]) and torch.equal(layer.values, orig_values[i]), \
+            f"layer {i} of the original state_cache was mutated"
+    print("[expand_state_cache] matches deepcopy+repeat_interleave reference, no aliasing")
+
+
 @pytest.mark.skipif(not os.environ.get("RUN_SLOW"), reason="downloads the real typical-small "
                      "checkpoint (Qwen3-1.7B-Base + LoRA); set RUN_SLOW=1 to run")
 def test_typical_matches_pcdm_decider():
@@ -218,6 +261,7 @@ def test_max_option_tokens_is_noop_when_nothing_truncated():
 
 
 if __name__ == "__main__":
+    test_expand_state_cache_matches_deepcopy_and_does_not_alias()
     if not os.environ.get("RUN_SLOW"):
         os.environ["RUN_SLOW"] = "1"
     test_typical_matches_pcdm_decider()

inference/typical/native.py

@@ -10,6 +10,8 @@ typical-small-preview) trains with nc_head n2/n2n3.
 """
 import bisect
 import copy
+import functools
+import json
 import string
 
 import torch
@@ -68,8 +70,55 @@ def _render_query_only(query: str, cand_texts: list[str]):
     return query + "\n", []
 
 
+# Verbatim from pcdm_jev.decider.SEMIF_SYSTEM / native.SEMIF_SYSTEM (SemIf's core.py DIRECT_SYSTEM).
+SEMIF_SYSTEM = ("Apply the supplied criterion to the supplied evidence. Choose exactly one listed option. "
+                "Respond with only its uppercase letter, with no explanation or reasoning.")
+
+
+@functools.lru_cache(maxsize=4)
+def _semif_head_tail(tok):
+    """(pre_head, tail): chat-template text around a sentinel "evidence" value -- pre_head =
+    system turn + user-turn opening + `{"evidence": "` (prepended to the state text to form the
+    KV-cached prefix), tail = end of user turn + assistant-turn opening (appended after the rest
+    of the JSON, see _render_semif). Verbatim port of native._semif_head_tail."""
+    sentinel = "\x00SEMIF_STATE\x00"
+    messages = [{"role": "system", "content": SEMIF_SYSTEM}, {"role": "user", "content": sentinel}]
+    try:
+        full = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
+    except TypeError:
+        full = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    before, after = full.split(sentinel)
+    return before + '{"evidence": "', after
+
+
+def _render_semif(query: str, cand_texts: list[str], tail_text: str):
+    """native_v4 (--nc_render semif): closes the "evidence" string opened in the cached prefix,
+    then criterion + lettered options (json.dumps-escaped) + tail_text. spans[k] = (start, end)
+    of option k's description text -- K spans, no null line. Verbatim port of native._render_semif."""
+    text = '", "criterion": ' + json.dumps(query) + ', "options": ['
+    spans = []
+    for i, c in enumerate(cand_texts):
+        if i:
+            text += ", "
+        text += '{"letter": "' + _label(i) + '", "description": "'
+        esc = json.dumps(c)[1:-1]
+        spans.append((len(text), len(text) + len(esc)))
+        text += esc + '"}'
+    text += ']}' + tail_text
+    return text, spans
+
+
 RENDERS = {"letters": _render, "tags": _render_tags, "letters_nonull": _render_letters_nonull,
-           "query_only": _render_query_only}
+           "query_only": _render_query_only, "semif": _render_semif}
+
+
+def _render_row(tok, render: str, query: str, cand_texts: list[str]):
+    """Dispatch RENDERS[render](query, cand_texts) for every render except "semif", which also
+    needs the tokenizer's chat-template tail text (_semif_head_tail). Verbatim port of
+    native._render_row."""
+    if render == "semif":
+        return _render_semif(query, cand_texts, _semif_head_tail(tok)[1])
+    return RENDERS[render](query, cand_texts)
 
 
 def _is_bern_row(cand_texts, meta=None):
@@ -263,20 +312,73 @@ def _pool_matrix(spans, offsets, base, Kmax, T, K=None):
     return pool / pool.sum(-1, keepdim=True).clamp_min(1.0)
 
 
-def _cache_batch_repeat_interleave(cache, repeats: int):
-    """cache.batch_repeat_interleave(repeats), generalised to Qwen3.5's hybrid cache -- see
-    native.py's copy for the long form (its Gated-DeltaNet linear-attention layers have no
-    batch_repeat_interleave of their own; Qwen3's cache is unaffected, every layer hits the
-    first branch)."""
+def _expand_state_cache(cache, repeats: int):
+    """Replaces the old `copy.deepcopy(state_cache)` + `_cache_batch_repeat_interleave`: a
+    new Cache whose K/V tensors are `repeats`-wide views of `cache`'s (batch=1) K/V via
+    `torch.Tensor.expand` -- stride-0 broadcast, zero bytes copied -- instead of a full deep
+    copy of the (up to 1024-token) prefix followed by a real repeat_interleave copy on top.
+    Safe because every attention layer's `.update()` (transformers' DynamicLayer et al.)
+    does `self.keys = torch.cat([self.keys, key_states], dim=-2)`: a rebind to a brand-new
+    tensor, never an in-place write to the old one, so the expanded view of the original
+    prefix is read (by torch.cat, which handles stride-0 inputs like any other tensor -- no
+    approximation, bit-identical to a materialised repeat) but never mutated. Only the outer
+    Cache object and each layer object are shallow-copied (cheap: Python attribute dicts,
+    not tensor storage).
+    ponytail: Gated-DeltaNet (Qwen3.5 hybrid cache) linear-attention layers have no K/V
+    tensors, only small conv_states/recurrent_states dicts (O(1) in sequence length) that
+    genuinely get mutated by the recurrent update -- those still get a real
+    repeat_interleave, just onto a copied dict so the write can't alias back into the
+    original state_cache. Only plain K/V (DynamicLayer-family) and this dict-state shape are
+    handled -- a future model with e.g. DynamicIndexedLayer's extra indexer_keys would need
+    its own expand branch here."""
+    new = copy.copy(cache)
+    new.layers = []
     for layer in cache.layers:
-        if hasattr(layer, "batch_repeat_interleave"):
-            layer.batch_repeat_interleave(repeats)
-            continue
-        for i in range(getattr(layer, "number_of_states", 1)):
-            if layer.is_conv_states_initialized[i]:
-                layer.conv_states[i] = layer.conv_states[i].repeat_interleave(repeats, dim=0)
-            if layer.is_recurrent_states_initialized[i]:
-                layer.recurrent_states[i] = layer.recurrent_states[i].repeat_interleave(repeats, dim=0)
+        nl = copy.copy(layer)
+        if hasattr(nl, "batch_repeat_interleave"):
+            if nl.keys is not None and nl.keys.numel() > 0:
+                nl.keys = nl.keys.expand(repeats, *nl.keys.shape[1:])
+                nl.values = nl.values.expand(repeats, *nl.values.shape[1:])
+        else:
+            nl.conv_states = dict(layer.conv_states)
+            nl.recurrent_states = dict(layer.recurrent_states)
+            for i in range(getattr(nl, "number_of_states", 1)):
+                if nl.is_conv_states_initialized[i]:
+                    nl.conv_states[i] = nl.conv_states[i].repeat_interleave(repeats, dim=0)
+                if nl.is_recurrent_states_initialized[i]:
+                    nl.recurrent_states[i] = nl.recurrent_states[i].repeat_interleave(repeats, dim=0)
+        new.layers.append(nl)
+    return new
+
+
+@functools.lru_cache(maxsize=64)
+def _causal_template(q_len: int, kv_len: int, device_str: str):
+    """[q_len, kv_len] boolean causal template (True = attend), cached per (q_len, kv_len,
+    device) -- a pure function of shape, so a repeat call at the same (prefix_len,
+    suffix_len) bucket reuses it instead of rebuilding the two arange()s + comparison."""
+    device = torch.device(device_str)
+    past_len = kv_len - q_len
+    q_idx = torch.arange(q_len, device=device).unsqueeze(1)
+    kv_idx = torch.arange(kv_len, device=device).unsqueeze(0)
+    return kv_idx <= (past_len + q_idx)
+
+
+@functools.lru_cache(maxsize=64)
+def _position_ids_row(T: int, Ls: int, device_str: str):
+    """[T] position ids (Ls, Ls+1, ..., Ls+T-1), cached per (T, Ls, device)."""
+    return torch.arange(T, device=torch.device(device_str)) + Ls
+
+
+@functools.lru_cache(maxsize=64)
+def _sink_ones(m: int, Ls: int, device_str: str):
+    """[m, Ls] all-ones long tensor (the cached prefix is always fully attended), cached per
+    (m, Ls, device) -- avoids a fresh CPU allocation on every decision."""
+    return torch.ones(m, Ls, dtype=torch.long, device=torch.device(device_str))
+
+
+@functools.lru_cache(maxsize=64)
+def _arange(n: int, device_str: str):
+    return torch.arange(n, device=torch.device(device_str))
 
 
 def _causal_pad_mask(attn2d, q_len, dtype):
@@ -284,10 +386,7 @@ def _causal_pad_mask(attn2d, q_len, dtype):
     past_key_values path mis-handles a right-padded batch sharing one KV cache; this
     reproduces the un-batched per-row computation exactly."""
     batch, kv_len = attn2d.shape
-    past_len = kv_len - q_len
-    q_idx = torch.arange(q_len, device=attn2d.device).unsqueeze(1)
-    kv_idx = torch.arange(kv_len, device=attn2d.device).unsqueeze(0)
-    causal = kv_idx <= (past_len + q_idx)
+    causal = _causal_template(q_len, kv_len, str(attn2d.device))
     allowed = causal.unsqueeze(0) & attn2d.bool().unsqueeze(1)
     mask = torch.zeros(batch, 1, q_len, kv_len, dtype=dtype, device=attn2d.device)
     return mask.masked_fill(~allowed.unsqueeze(1), torch.finfo(dtype).min)
@@ -310,20 +409,30 @@ def _truncate_options(tok, cand_texts: list[str], max_tokens: int | None) -> lis
 
 
 @torch.inference_mode()
-def encode_state(head, model, state: str, max_state: int = 256):
-    """[eos] + state -> (past_key_values, Ls), the once-per-state prefix encode that used to
+def encode_state(head, model, state: str, max_state: int = 256, semif: bool | None = None):
+    """[eos] + state (or, --nc_render semif, the SemIf chat-template wrapper + state -- see
+    _semif_head_tail) -> (past_key_values, Ls), the once-per-state prefix encode that used to
     run inside native_kv_decide on every call. Split out so a caller (Typical's state LRU)
     can compute it once and reuse the KV cache across many decisions on the same state --
     see native_kv_decide's `state_cache` param. Pure function of (backbone weights, state
-    text): deterministic in eval/inference_mode, so a cached entry is bit-identical to a
+    text, semif): deterministic in eval/inference_mode, so a cached entry is bit-identical to a
     fresh encode, not just numerically close. @inference_mode here (not just relying on the
-    caller) because the resulting KV tensors get copy.deepcopy'd later (native_kv_decide's
-    per-chunk batch_repeat_interleave) -- a non-leaf, grad-tracked tensor can't be
-    deepcopy'd, so this must never run under regular autograd tracking, regardless of what
-    context the caller happens to be in."""
+    caller) because the resulting KV tensors get `.expand()`'d and read (never in-place
+    mutated) by every subsequent decision on this state (see `_expand_state_cache`) -- an
+    inference_mode tensor is a stable, non-grad-tracked view source, so this must never run
+    under regular autograd tracking, regardless of what context the caller happens to be in.
+    semif=None defaults to `model.render == "semif"`; native_kv_decide passes semif=False
+    explicitly for a bern (yes/no) row's own prefix, since those always use the plain eos+state
+    sink regardless of nc_render (see native.py's module docstring)."""
+    if semif is None:
+        semif = getattr(model, "render", "letters") == "semif"
     tok, dev, lm = head.backbone.tokenizer, head.device, head.backbone.model
     with record_function("typical/tokenize_state"):
-        prefix = torch.tensor([[tok.eos_token_id] + _ids(tok, [state], max_state)[0]], device=dev)
+        if semif:
+            prefix_ids = _ids(tok, [_semif_head_tail(tok)[0] + state], max_state)[0]
+        else:
+            prefix_ids = [tok.eos_token_id] + _ids(tok, [state], max_state)[0]
+        prefix = torch.tensor([prefix_ids], device=dev)
     with record_function("typical/state_prefix_forward"):
         cache = lm(input_ids=prefix, use_cache=True).past_key_values
     return cache, prefix.shape[1]
@@ -333,17 +442,25 @@ def encode_state(head, model, state: str, max_state: int = 256):
 def native_kv_decide(head, model, state, queries, chunk: int = 32, max_state: int = 256,
                      max_suffix: int = MAX_SUFFIX, vec_cache=None, state_cache=None,
                      max_option_tokens: int | None = DEFAULT_MAX_OPTION_TOKENS):
-    """Serving path: [eos] + state encoded ONCE into a KV cache on head.backbone.model; every
-    query's suffix runs against that cached prefix, `chunk` queries per forward. queries =
-    [(query, candidates)] -> list of [K_i + 1] probability vectors (null last). Verbatim port
-    of native.native_kv_decide (head/model contract unchanged -- head needs only .backbone
-    (tokenizer + model) and .device; model is a NativeHead), plus two additions: `state_cache`,
-    an optional (past_key_values, Ls) pair from encode_state -- pass it to skip re-encoding
-    the state prefix (Typical.choice/score/noul/decide do this automatically via an LRU; see
-    core.py) -- and `max_option_tokens` ("suffix diet": caps each rendered option's length,
-    see _truncate_options; None/0 disables, matching the training repo's untruncated
-    behaviour exactly). No hierarchical fallback: a suffix longer than max_suffix raises
-    (raise max_suffix instead)."""
+    """Serving path: [eos] + state (or, --nc_render semif, the SemIf chat-template wrapper +
+    state) encoded ONCE into a KV cache on head.backbone.model; every query's suffix runs
+    against that cached prefix, `chunk` queries per forward. queries = [(query, candidates)] ->
+    list of [K_i + 1] probability vectors (null last). Verbatim port of native.native_kv_decide
+    (head/model contract unchanged -- head needs only .backbone (tokenizer + model) and .device;
+    model is a NativeHead), plus two additions: `state_cache`, an optional (past_key_values, Ls)
+    pair from encode_state -- pass it to skip re-encoding the state prefix (Typical.choice/
+    score/noul/decide do this automatically via an LRU; see core.py) -- and `max_option_tokens`
+    ("suffix diet": caps each rendered option's length, see _truncate_options; None/0 disables,
+    matching the training repo's untruncated behaviour exactly). No hierarchical fallback: a
+    suffix longer than max_suffix raises (raise max_suffix instead).
+
+    --nc_render semif: a bern (yes/no) row always renders query_only against the PLAIN eos+state
+    prefix, never the semif JSON wrapper (see native.py's module docstring), so it can't share
+    `state_cache` when one is supplied (Typical's per-state LRU key/encode_state default to
+    building the semif-wrapped cache whenever model.render=="semif" -- see encode_state). A
+    call that mixes bern and non-bern queries under a semif checkpoint therefore builds a
+    second, plain prefix on demand for the bern rows only; every other render (incl. no bern
+    rows present) is the original single-cache loop, unchanged."""
     if model.use2:
         raise NotImplementedError(
             "nc_head n2/n2n3 (Qwen3-Embedding candidate vectors) isn't ported to the "
@@ -352,47 +469,63 @@ def native_kv_decide(head, model, state, queries, chunk: int = 32, max_state: in
             "checkpoint requires it."
         )
     tok, dev, lm = head.backbone.tokenizer, head.device, head.backbone.model
-    state_cache, Ls = state_cache if state_cache is not None else encode_state(head, model, state, max_state)
-    out = []
-    for start in range(0, len(queries), chunk):
-        qs = queries[start:start + chunk]
-        m = len(qs)
-        with record_function("typical/tokenize_suffix"):
-            # is_bern/yes-idx/cmask all key off candidate COUNT and the literal "yes"/"no"
-            # text (never truncated -- 1 token), so they read the original `c`; only the
-            # rendered text is diet'd.
-            is_bern = [model.noul_head == "bern" and _is_bern_row(c) for _, c in qs]
-            rendered = [RENDERS["query_only" if b else getattr(model, "render", "letters")]
-                        (q, _truncate_options(tok, c, max_option_tokens))
-                        for b, (q, c) in zip(is_bern, qs)]
-            enc = tok([r[0] for r in rendered], add_special_tokens=False, return_offsets_mapping=True)
-            x_ids, offs = enc["input_ids"], enc["offset_mapping"]
-            assert max(len(x) for x in x_ids) <= max_suffix, f"suffix > max_suffix={max_suffix}; raise it"
-            input_ids, am, lengths = _pack(tok, [[]] * m, x_ids, tail=[tok.eos_token_id], sink=False)
-        with record_function("typical/mask_build"):
-            T = input_ids.shape[1]
-            Kmax = max(len(c) for _, c in qs)
-            pool = torch.stack([_pool_matrix(r[1], o, 0, Kmax, T, K=len(c)) for r, o, (_, c) in zip(rendered, offs, qs)])
-            cmask = torch.zeros(m, Kmax, dtype=torch.bool)
-            for i, (_, c) in enumerate(qs):
-                cmask[i, :len(c)] = True
-
-        with record_function("typical/state_cache_deepcopy"):
-            cache = copy.deepcopy(state_cache)
-            _cache_batch_repeat_interleave(cache, m)
-        with record_function("typical/suffix_forward"):
-            attn = torch.cat([torch.ones(m, Ls, dtype=torch.long), am], dim=1).to(dev)
-            position_ids = (torch.arange(T) + Ls).expand(m, -1).to(dev)
-            H = lm(input_ids=input_ids.to(dev), attention_mask=_causal_pad_mask(attn, T, lm.dtype),
-                   position_ids=position_ids, past_key_values=cache, use_cache=False).last_hidden_state.float()
-        with record_function("typical/head_forward"):
-            h = H[torch.arange(m, device=dev), (lengths - 1).to(dev)]
-            C3 = torch.bmm(pool.to(dev), H)
-            C2 = None  # n2 unsupported -- guarded above
-            yi = _yes_idx([c for _, c in qs], dev) if model.noul_head == "bern" else None
-            bm = torch.tensor(is_bern, dtype=torch.bool, device=dev) if model.noul_head == "bern" else None
-            probs = torch.softmax(model(h, cmask.to(dev), C3, C2, yes_idx=yi, bern_mask=bm), dim=-1)
-        with record_function("typical/cpu_sync"):
-            for i, (_, c) in enumerate(qs):
-                out.append(torch.cat([probs[i, :len(c)], probs[i, -1:]]))
+    base_render = getattr(model, "render", "letters")
+    # is_bern/yes-idx/cmask all key off candidate COUNT and the literal "yes"/"no" text (never
+    # truncated -- 1 token), so they read the original `c`; only the rendered text is diet'd.
+    is_bern = [model.noul_head == "bern" and _is_bern_row(c) for _, c in queries]
+    renders = ["query_only" if b else base_render for b in is_bern]
+    out = [None] * len(queries)
+
+    def run_group(idxs, cache, Ls, tail_ids):
+        if not idxs:
+            return
+        for start in range(0, len(idxs), chunk):
+            grp = idxs[start:start + chunk]
+            qs = [queries[i] for i in grp]
+            rnds = [renders[i] for i in grp]
+            m = len(qs)
+            with record_function("typical/tokenize_suffix"):
+                rendered = [_render_row(tok, r, q, _truncate_options(tok, c, max_option_tokens))
+                            for r, (q, c) in zip(rnds, qs)]
+                enc = tok([r[0] for r in rendered], add_special_tokens=False, return_offsets_mapping=True)
+                x_ids, offs = enc["input_ids"], enc["offset_mapping"]
+                assert max(len(x) for x in x_ids) <= max_suffix, f"suffix > max_suffix={max_suffix}; raise it"
+                input_ids, am, lengths = _pack(tok, [[]] * m, x_ids, tail=tail_ids, sink=False)
+            with record_function("typical/mask_build"):
+                T = input_ids.shape[1]
+                Kmax = max(len(c) for _, c in qs)
+                pool = torch.stack([_pool_matrix(r[1], o, 0, Kmax, T, K=len(c)) for r, o, (_, c) in zip(rendered, offs, qs)])
+                cmask = torch.zeros(m, Kmax, dtype=torch.bool)
+                for i, (_, c) in enumerate(qs):
+                    cmask[i, :len(c)] = True
+
+            with record_function("typical/state_cache_expand"):
+                exp_cache = _expand_state_cache(cache, m)
+            with record_function("typical/suffix_forward"):
+                attn = torch.cat([_sink_ones(m, Ls, "cpu"), am], dim=1).to(dev)
+                position_ids = _position_ids_row(T, Ls, str(dev)).expand(m, -1)
+                H = lm(input_ids=input_ids.to(dev), attention_mask=_causal_pad_mask(attn, T, lm.dtype),
+                       position_ids=position_ids, past_key_values=exp_cache, use_cache=False).last_hidden_state.float()
+            with record_function("typical/head_forward"):
+                h = H[_arange(m, str(dev)), (lengths - 1).to(dev)]
+                C3 = torch.bmm(pool.to(dev), H)
+                C2 = None  # n2 unsupported -- guarded above
+                yi = _yes_idx([c for _, c in qs], dev) if model.noul_head == "bern" else None
+                bm = (torch.tensor([r == "query_only" for r in rnds], dtype=torch.bool, device=dev)
+                      if model.noul_head == "bern" else None)
+                probs = torch.softmax(model(h, cmask.to(dev), C3, C2, yes_idx=yi, bern_mask=bm), dim=-1)
+            with record_function("typical/cpu_sync"):
+                for row, i in enumerate(grp):
+                    _, c = qs[row]
+                    out[i] = torch.cat([probs[row, :len(c)], probs[row, -1:]])
+
+    non_bern, bern = [i for i, b in enumerate(is_bern) if not b], [i for i, b in enumerate(is_bern) if b]
+    if base_render != "semif" or not bern:
+        cache, Ls = state_cache if state_cache is not None else encode_state(head, model, state, max_state)
+        run_group(list(range(len(queries))), cache, Ls, [tok.eos_token_id])
+    else:
+        semif_cache, semif_Ls = state_cache if state_cache is not None else encode_state(head, model, state, max_state)
+        run_group(non_bern, semif_cache, semif_Ls, [])
+        plain_cache, plain_Ls = encode_state(head, model, state, max_state, semif=False)
+        run_group(bern, plain_cache, plain_Ls, [tok.eos_token_id])
     return out