Increase adapter reuse + consistent reuse pattern

hf_adapters/_dspark_common.py

@@ -399,7 +399,7 @@ def install_spyre_markov(model):
     correction runs on host with only the ``w2`` matmul on-device — no device mix.
     No-op when the drafter has no markov head.
     """
-    from deepspec.utils.sampling import sample_tokens
+    from deepspec.utils.sampling import sample_tokens  # type: ignore[import-not-found]
 
     mh = getattr(model, "markov_head", None)
     if mh is None or not hasattr(mh, "markov_w2"):

hf_adapters/hf_dspark_gemma4.py

@@ -46,12 +46,10 @@
     _pad_markov_w2,
     install_spyre_compute_logits,
     install_spyre_markov,
+    run_draft_block,
     snapshot_cpu_embeddings,
     snapshot_cpu_fc,
 )
-from hf_adapters._dspark_common import (  # noqa: F401  (re-exported as the public forward)
-    run_draft_block as _run_draft_block,
-)
 from hf_adapters.hf_common import (
     InvFreqShim,
     PrecomputedRotaryEmbedding,
@@ -60,6 +58,7 @@
 )
 from hf_adapters.hf_gemma4 import _patch_gemma4_rmsnorm
 
+_run_draft_block = run_draft_block  # reuse the common runner
 CTX_PAD = 56
 
 
@@ -132,7 +131,9 @@ def block_forward(hidden_states, target_hidden_states, selected_freqs, attn_mask
 
 def prepare_for_spyre(model):
     """Apply Spyre adaptations to the Gemma4 DSpark drafter in-place."""
-    from deepspec.modeling.dspark.gemma4.modeling import Gemma4RMSNorm
+    from deepspec.modeling.dspark.gemma4.modeling import (  # type: ignore[import-not-found]
+        Gemma4RMSNorm,
+    )
 
     block_size = int(model.block_size)
     kv_pad = ((CTX_PAD + block_size + 31) // 32) * 32

hf_adapters/hf_dspark_granite.py

@@ -34,13 +34,9 @@
 Usage: see ``hf_dspark_qwen3``.
 """
 
-# The block-propose forward is shared across families — reuse the common runner
-# under the adapter's public ``_run_draft_block`` name (see hf_dspark_qwen3).
-from hf_adapters._dspark_common import prepare_dspark_common
-from hf_adapters._dspark_common import (  # noqa: F401  (re-exported as the public forward)
-    run_draft_block as _run_draft_block,
-)
+from hf_adapters._dspark_common import prepare_dspark_common, run_draft_block
 
+_run_draft_block = run_draft_block  # reuse the common runner
 CTX_PAD = 56
 
 

hf_adapters/hf_dspark_qwen3.py

@@ -37,22 +37,19 @@
         model, draft_input_ids, target_hidden_states, selected_freqs, ctx_valid_len)
 """
 
-# The block-propose forward is identical across the drafter families — reuse the
-# shared runner directly (exposed under the adapter's public ``_run_draft_block``
-# name that ``resolve_adapter_module``/tests look up).
-from hf_adapters._dspark_common import prepare_dspark_common
-from hf_adapters._dspark_common import (  # noqa: F401  (re-exported as the public forward)
-    run_draft_block as _run_draft_block,
-)
+from hf_adapters._dspark_common import prepare_dspark_common, run_draft_block
 
+_run_draft_block = run_draft_block  # reuse the common runner
 # Fixed context / kv widths (stick-aligned). CTX_PAD=56 keeps kv_pad =
 # round32(CTX_PAD + block_size=7) = 64 — the proven-compilable attention width.
 CTX_PAD = 56
 
 
 def prepare_for_spyre(model):
     """Apply Spyre adaptations to the Qwen3 DSpark drafter in-place."""
-    from deepspec.modeling.dspark.qwen3.modeling import Qwen3RMSNorm
+    from deepspec.modeling.dspark.qwen3.modeling import (  # type: ignore[import-not-found]
+        Qwen3RMSNorm,
+    )
 
     block_size = int(model.block_size)
     kv_pad = ((CTX_PAD + block_size + 31) // 32) * 32

hf_adapters/hf_granite.py

@@ -32,6 +32,7 @@
     pad_lm_head,
     patch_rmsnorm,
     prepare_rope_and_heads,
+    text_config,
 )
 
 
@@ -93,8 +94,7 @@ def _run_forward(
         cache_position,
     )
     logits = model.lm_head(h)
-    logits = logits / model.config.logits_scaling
-    return logits
+    return logits / text_config(model.config).logits_scaling
 
 
 def prepare_for_spyre(model):

hf_adapters/hf_granite_vision.py

@@ -41,15 +41,17 @@
 
 import torch
 
+from hf_adapters import hf_granite
 from hf_adapters.hf_common import (
     get_backbone,
     make_standard_gqa_block,
     pad_lm_head,
     patch_rmsnorm,
     prepare_rope_and_heads,
-    text_config,
 )
-from hf_adapters.hf_granite import _run_backbone_forward
+
+_run_backbone_forward = hf_granite._run_backbone_forward
+_run_forward = hf_granite._run_forward
 
 
 def load_hf_model(model_path, dtype=torch.float16):
@@ -96,34 +98,3 @@ def prepare_for_spyre(model):
     model._spyre_compiled_blocks = [
         make_standard_gqa_block(layer, True) for layer in get_backbone(model).layers
     ]
-
-
-def _run_forward(
-    model,
-    input_ids,
-    position_ids,
-    attn_mask,
-    key_caches,
-    value_caches,
-    is_filling,
-    token_index,
-    cache_position,
-):
-    """Granite Vision text causal-LM forward: backbone + head / scaling.
-
-    Identical to ``hf_granite._run_forward`` except ``logits_scaling`` lives on
-    the nested ``text_config`` rather than the top-level VLM config.
-    """
-    h = _run_backbone_forward(
-        model,
-        input_ids,
-        position_ids,
-        attn_mask,
-        key_caches,
-        value_caches,
-        is_filling,
-        token_index,
-        cache_position,
-    )
-    logits = model.lm_head(h)
-    return logits / text_config(model.config).logits_scaling

hf_adapters/hf_granitemoehybrid.py

@@ -44,6 +44,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from hf_adapters import hf_granite
 from hf_adapters.hf_common import (
     apply_rope_matmul,
     get_backbone,
@@ -54,6 +55,9 @@
     split_fused_linear,
 )
 
+_run_backbone_forward = hf_granite._run_backbone_forward
+_run_forward = hf_granite._run_forward
+
 
 def _make_compiled_block(layer, res_mult, gate_proj, up_proj):
     """Compiled block for Granite 4.0 dense: split MLP, multipliers."""
@@ -119,68 +123,6 @@ def block_forward(
     return torch.compile(block_forward, dynamic=False)
 
 
-def _run_backbone_forward(
-    model,
-    input_ids,
-    position_ids,
-    attn_mask,
-    key_caches,
-    value_caches,
-    is_filling,
-    token_index,
-    cache_position,
-):
-    """Granite 4.0 backbone: embedding * multiplier, blocks, norm."""
-    backbone = get_backbone(model)
-    h = backbone.embed_tokens(input_ids)
-    h = h * model.config.embedding_multiplier
-
-    selected_freqs = model._spyre_rope(h, position_ids)
-
-    for i, compiled_block in enumerate(model._spyre_compiled_blocks):
-        h, key_caches[i], value_caches[i] = compiled_block(
-            h,
-            selected_freqs,
-            attn_mask,
-            key_caches[i],
-            value_caches[i],
-            is_filling,
-            token_index,
-            cache_position,
-        )
-
-    h = backbone.norm(h)
-    return h
-
-
-def _run_forward(
-    model,
-    input_ids,
-    position_ids,
-    attn_mask,
-    key_caches,
-    value_caches,
-    is_filling,
-    token_index,
-    cache_position,
-):
-    """Granite 4.0 causal-LM forward: backbone + head / scaling."""
-    h = _run_backbone_forward(
-        model,
-        input_ids,
-        position_ids,
-        attn_mask,
-        key_caches,
-        value_caches,
-        is_filling,
-        token_index,
-        cache_position,
-    )
-    logits = model.lm_head(h)
-    logits = logits / model.config.logits_scaling
-    return logits
-
-
 def prepare_for_spyre(model):
     """Apply Spyre adaptations to Granite 4.0 dense model in-place."""
     from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (

hf_adapters/hf_mistral3.py

@@ -49,14 +49,14 @@
     outputs = model.generate(tokenizer, ["Hello!"], max_new_tokens=32)
 """
 
+from hf_adapters import hf_mistral
 from hf_adapters.hf_common import (
     get_backbone,
     prepare_standard_gqa,
 )
-from hf_adapters.hf_mistral import (
-    _run_backbone_forward,  # noqa: F401  re-exported as adapter module API
-    _run_forward,  # noqa: F401  re-exported as adapter module API
-)
+
+_run_backbone_forward = hf_mistral._run_backbone_forward
+_run_forward = hf_mistral._run_forward
 
 
 def load_hf_model(model_path, dtype):

hf_adapters/hf_xlm_roberta.py

@@ -15,7 +15,7 @@
 """
 HuggingFace Transformers adapter for XLM-RoBERTa encoder-only models on Spyre.
 
-Supports models with XLMRobertaConfig (e.g. BAAI/bge-m3,
+Supports models with RoBERTa/XLMRobertaConfig (e.g. BAAI/bge-m3,
 intfloat/multilingual-e5-large, sentence-transformers/paraphrase-multilingual-*).
 
 Structurally identical to BERT (same attention/FFN/post-LN module names) so the
@@ -27,47 +27,26 @@
   ``padding_idx``. ``prefill_encoder`` synthesizes 0-based position ids that are
   correct for BERT but wrong for XLM-R, so this adapter overrides
   ``_run_backbone_forward`` to recompute position ids from ``input_ids`` itself.
-- Embedding sum order differs (``word + position`` then ``+ token_type`` vs
-  BERT's three-way add). Mathematically equivalent — same final tensor. The
-  token-type add goes through ``add_token_type_embedding``, which broadcasts the
-  single row directly when ``type_vocab_size == 1`` (as in bge-m3 /
-  bge-reranker-v2-m3) because a one-row gather does not lower on Spyre.
 """
 
 from hf_adapters.hf_bert import _make_compiled_encoder_block
 from hf_adapters.hf_common import (
     BLOCK_SIZE,
-    add_token_type_embedding,
+    encoder_backbone_forward,
     fairseq_position_ids,
     get_backbone,
     pad_attention_heads_simple,
 )
 
 
 def _run_backbone_forward(model, input_ids, attn_mask, position_ids, token_type_ids):
-    """Encoder backbone forward with XLM-R position ids.
-
-    Modified version of ``encoder_backbone_forward``. To maintain the signature
-    of the original function, we pass in ``position_ids`` as an argument, but
-    compute the XLM-R-style positions from ``input_ids``. Otherwise
-    follows ``encoder_backbone_forward``: word + position + token_type embed,
-    LayerNorm, then the compiled encoder blocks with the Spyre layout-fixup
-    clones around each block.
-    """
-    backbone = get_backbone(model)
-    emb = backbone.embeddings
-
+    """Encoder backbone forward with RoBERTa/XLM-R position ids."""
+    emb = get_backbone(model).embeddings
     pos_ids = fairseq_position_ids(input_ids, emb.padding_idx)
 
-    h = emb.word_embeddings(input_ids) + emb.position_embeddings(pos_ids)
-    h = add_token_type_embedding(h, emb, token_type_ids)
-    h = emb.LayerNorm(h)
-    h = h.clone() if h.device.type == "spyre" else h
-    for compiled_block in model._spyre_compiled_blocks:
-        h = compiled_block(h, attn_mask)
-        if h.device.type == "spyre":
-            h = h.clone()
-    return h
+    return encoder_backbone_forward(
+        model, input_ids, attn_mask, pos_ids, token_type_ids
+    )
 
 
 _is_encoder_only = True

tests/spyre/test_vlm_e2e_spyre.py

@@ -146,7 +146,6 @@ def _adapter_teacher_forced_steps(
 
     model_d_type = get_model_dtype(model)
     backbone = adapter.get_backbone(model)
-    # emb_mult = backbone.embedding_multiplier
     # Falls back to 1.0 for models (Mistral) that don't scale embeddings
     emb_mult = getattr(backbone, "embedding_multiplier", 1.0)