update lm head last token pruning - #288
Merged
Merged
Conversation
Contributor
Author
|
@copilot merge main branch into this PR. |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds an opt-in prune_lm_head flag that, when enabled, prunes the LM head computation to only the last token (prefill optimization) by inserting a Gather before the LM head. It also introduces an export_package() helper in the ORT-GenAI integration to produce a fully loadable onnxruntime-genai directory (ONNX + config artifacts) from an already-built ModelPackage.
Changes:
- Add
flags.prune_lm_head(env:MOBIUS_PRUNE_LM_HEAD) and wire it into the baseCausalLMModel.forward()path. - Add unit tests for the pruning flag behavior in the base model build path.
- Add
export_package()to the ORT-GenAI integration, refactorauto_export()to delegate to it, and add tests/docs for the new helper.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/models/base.py | Adds last-token hidden-state selection before lm_head when flags.prune_lm_head is enabled. |
| src/mobius/models/_models_test.py | Adds tests covering default vs pruned logits shape behavior for CausalLMModel. |
| src/mobius/integrations/ort_genai/auto_export.py | Adds export_package() and refactors auto_export() to use it; updates module docstring. |
| src/mobius/integrations/ort_genai/auto_export_test.py | Adds unit tests for export_package() behavior and argument propagation. |
| src/mobius/integrations/ort_genai/init.py | Exposes export_package in the integration public API. |
| src/mobius/_model_package.py | Updates ModelPackage.save() docstring to point users to export_package() for ORT-GenAI directories. |
| src/mobius/_flags.py | Adds the prune_lm_head flag, env var support, and documentation. |
…-last-token-pruning Signed-off-by: GitHub <noreply@github.com> # Conflicts: # src/mobius/_flags.py # src/mobius/integrations/ort_genai/auto_export.py # src/mobius/integrations/ort_genai/auto_export_test.py Co-authored-by: rui-ren <15321482+rui-ren@users.noreply.github.com>
Contributor
Done. Merged |
rui-ren
changed the title
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Performance Comparison
|
|
The author of this PR, rui-ren, is not an activated member of this organization on Codecov. |
Member
|
Maybe we should consider creating an option like we do for use_static_cache, as the feature flags are designed to be temporary and would eventually be removed when the feature stabilizes. Do you think it would make sense to create a new task for this? Or maybe create an option to an existing task? |
Contributor
Author
ok, sure, make sense to create a new task for this. |
Member
|
@copilot please address comments and update from main |
Remove the MOBIUS_PRUNE_LM_HEAD feature flag and replace it with a prune_lm_head constructor parameter on CausalLMTask (and HybridCausalLMTask), mirroring the stable-API pattern already used by the static_cache option. When prune_lm_head=True, the task inserts Gather(axis=1, index=-1) followed by Unsqueeze(axis=1) on the logits output after calling module(), changing the output shape from [B, S, vocab] to [B, 1, vocab]. ONNX Runtime's graph optimizer can push this Gather backward through the LM head MatMul to avoid computing all-token logits during prefill. Update tests to use CausalLMTask(prune_lm_head=True) instead of the override_flags context manager. Also merges main branch changes (tencent_q1_0_use_native_2bit and static_cache_bias flags). Signed-off-by: Copilot <copilot@github.com> Signed-off-by: GitHub <noreply@github.com>
|
|
Contributor
Done in commit Moved # Before (global flag)
with override_flags(prune_lm_head=True):
pkg = build(model_id)
# After (task option)
pkg = build(model_id, task=CausalLMTask(prune_lm_head=True))The |
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (5)
src/mobius/tasks/_causal_lm.py:332
- Same as above:
Gather(..., idx=-1)depends on negative indices. Derive the last index fromShape(logits)to avoid relying on undefined/implementation-specific behavior.
last_idx = op.Constant(value_int=-1) # scalar (rank-0) INT64
last_token_logits = op.Gather(logits, last_idx, axis=1) # [B, vocab]
logits = op.Unsqueeze(last_token_logits, op.Constant(value_ints=[1])) # [B, 1, vocab]
src/mobius/tasks/_causal_lm.py:87
- The docstring says pruning inserts a Gather "before the LM head" so only the last hidden state is projected, but the implementation below prunes after logits are computed (Gather on
[B,S,vocab]). This is misleading and also overstates the guaranteed perf win (it relies on ORT pushing the Gather backward through the projection).
prune_lm_head: If ``True``, insert ``Gather(axis=1, index=-1)``
before the LM head so only the last token's hidden state is
projected to logits. Output logits shape becomes ``[B, 1,
vocab]`` instead of ``[B, S, vocab]``, reducing prefill cost
for large-vocabulary models. Set this when the downstream
src/mobius/tasks/_causal_lm.py:215
- Using
Gather(..., idx=-1)relies on negative indices, which are not guaranteed by the ONNX Gather spec across runtimes. Compute the last index fromShape(logits)instead to keep the graph spec-compliant and portable.
This issue also appears on line 330 of the same file.
last_idx = op.Constant(value_int=-1) # scalar (rank-0) INT64
last_token_logits = op.Gather(logits, last_idx, axis=1) # [B, vocab]
logits = op.Unsqueeze(last_token_logits, op.Constant(value_ints=[1])) # [B, 1, vocab]
src/mobius/tasks/_causal_lm.py:99
- The PR description says this is enabled via
MOBIUS_PRUNE_LM_HEAD=1orflags.prune_lm_head=True, but the implementation only supports passingprune_lm_headwhen constructingCausalLMTask/HybridCausalLMTask. If the env-var / global flags integration is a requirement, it’s currently missing.
def __init__(
self,
*,
static_cache: bool = False,
max_seq_len: int | None = None,
prune_lm_head: bool = False,
):
self._static_cache = static_cache
self._max_seq_len = max_seq_len
self._prune_lm_head = prune_lm_head
src/mobius/models/_models_test.py:233
test_prune_emits_gather_on_logitsonly checks the inferred output shape. That can miss regressions where shape inference still yields[B,1,V]but the pruning ops are not actually wired (or get optimized away unexpectedly). Making the test assert theUnsqueeze <- Gatherproducer chain keeps it directly tied to the intended graph edit.
logits = next(v for v in model.graph.outputs if v.name == "logits")
# Logits must still be rank-3 [B, 1, vocab] (NOT rank-4 [B, 1, 1, V])
assert len(logits.shape) == 3, (
f"Expected rank-3 logits [B, 1, V] after pruning, got rank "
f"{len(logits.shape)}: shape={list(logits.shape)!r}"
)
# Pruned: dim 1 must be the literal integer 1
seq_dim = logits.shape[1]
assert seq_dim == 1, f"Expected logits dim 1 to be 1 after pruning, got {seq_dim!r}"
# Last dim is still the vocabulary size
config = make_config()
assert logits.shape[2] == config.vocab_size
justinchuby
marked this pull request as ready for review
July 30, 2026 05:33
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Member
|
@copilot update from main and make this a feature |
…-last-token-pruning Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com> Signed-off-by: GitHub <noreply@github.com>
Add prune-lm-head to the cargo-style build features and thread it through build()/build_from_module() so CLI and programmatic builds share the same task-controlled pruning path. Cover the new feature wiring in CLI and builder tests. Signed-off-by: GitHub <noreply@github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.
Suppressed comments (4)
src/mobius/_builder.py:242
- PR description says this is also controllable via
MOBIUS_PRUNE_LM_HEAD=1/flags.prune_lm_head = True, but the implementation only wires the new build() / build_from_module() parameter + CLI feature; there is noprune_lm_headruntime flag in src/mobius/_flags.py and build() doesn’t consultflagsfor this setting. Either update the description or add a runtime flag and have prune_lm_head default to it.
if prune_lm_head:
task = _enable_pruned_lm_head_task(task)
resolved_task = get_task(task)
src/mobius/_builder.py:128
- _enable_pruned_lm_head_task() reconstructs CausalLMTask by reading private attributes and silently drops any CausalLMTask/HybridCausalLMTask subclasses (isinstance() matches but the returned task is the base class). This can break callers that pass a customized task instance. Prefer cloning the passed task and toggling the prune flag without mutating the caller’s instance, so subclasses/extra fields are preserved.
This issue also appears on line 240 of the same file.
def _enable_pruned_lm_head_task(task: str | ModelTask) -> str | ModelTask:
"""Return a task equivalent to *task* with LM-head pruning enabled."""
from mobius.tasks import CausalLMTask, HybridCausalLMTask
if task == "text-generation":
return CausalLMTask(prune_lm_head=True)
if task == "hybrid-text-generation":
return HybridCausalLMTask(prune_lm_head=True)
if isinstance(task, CausalLMTask):
return CausalLMTask(
static_cache=getattr(task, "_static_cache", False),
max_seq_len=getattr(task, "_max_seq_len", None),
prune_lm_head=True,
)
if isinstance(task, HybridCausalLMTask):
return HybridCausalLMTask(prune_lm_head=True)
raise ValueError(
"prune_lm_head=True is only supported for text-generation and "
"hybrid-text-generation tasks."
)
src/mobius/tasks/_causal_lm.py:87
- CausalLMTask docstring says pruning inserts Gather before the LM head and projects only the last token’s hidden state, but the implementation gathers on the logits after the module returns. The docs should match the actual graph transformation (and note that any compute savings rely on ORT pushing the Gather backwards).
prune_lm_head: If ``True``, insert ``Gather(axis=1, index=-1)``
before the LM head so only the last token's hidden state is
projected to logits. Output logits shape becomes ``[B, 1,
vocab]`` instead of ``[B, S, vocab]``, reducing prefill cost
for large-vocabulary models. Set this when the downstream
docs/cli_reference.md:188
- The
prune-lm-headdescription implies runtimes can always skip the full prefill LM-head projection, but the current implementation prunes logits after the LM head and relies on ORT graph optimizations to push the Gather backward. Wording should reflect that this is optimizer-dependent.
| `prune-lm-head` | Emit only final-token logits (`[B, 1, vocab]`) for supported causal-LM tasks so runtimes can skip full prefill LM-head projection. Use only when the downstream workflow does not need per-token logits. |
…-last-token-pruning # Conflicts: # src/mobius/models/_models_test.py
Keep prune-lm-head task-controlled while selecting the final hidden-state position before the vocabulary projection. Fail explicitly for custom forwards that do not support the feature and document the cargo-style build option. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 88ec109e-af0c-46c3-9806-731bc2b7342f Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated no new comments.
Suppressed comments (2)
src/mobius/_builder.py:124
_enable_pruned_lm_head_task()usesisinstance(task, CausalLMTask)/isinstance(task, HybridCausalLMTask), which will also match subclasses. That silently replaces a custom task subclass with a new baseCausalLMTask(...), potentially dropping overridden behavior (and violating the intent of failing explicitly for unsupported tasks). Prefer restricting this branch to the exact task types (or explicitly handling subclasses) to avoid accidental behavior changes.
if task == "text-generation":
return CausalLMTask(prune_lm_head=True)
if task == "hybrid-text-generation":
return HybridCausalLMTask(prune_lm_head=True)
if isinstance(task, CausalLMTask):
return CausalLMTask(
static_cache=getattr(task, "_static_cache", False),
max_seq_len=getattr(task, "_max_seq_len", None),
prune_lm_head=True,
)
if isinstance(task, HybridCausalLMTask):
return HybridCausalLMTask(prune_lm_head=True)
src/mobius/models/_models_test.py:200
- The class docstring says pruning inserts
Gather + Unsqueezeafter the LM head, but the implementation/test assertion chain shows it happens before the LM-head projection (on hidden states). This is misleading documentation for the new feature.
"""Tests for the ``prune_lm_head`` option in :class:`CausalLMTask`.
When ``True``, a ``Gather + Unsqueeze`` is inserted after the LM head
so logits are produced for only the last sequence position.
Mirrors onnxruntime-genai Model Builder's ``prune_lm_head`` opt-in.
justinchuby
approved these changes
Aug 4, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR Description
Adds opt-in final-token LM-head pruning as the cargo-style build feature:
The same behavior is available programmatically through
build(prune_lm_head=True)andbuild_from_module(..., prune_lm_head=True).When enabled, Mobius selects the final hidden-state position before the LM-head projection, changing logits from
[B, S, vocab]to[B, 1, vocab]. This avoids computing unused per-token vocabulary logits during prefill.The default remains disabled to preserve workflows requiring per-token logits, including logprob scoring, speculative decoding, and multi-token generation. Models using unsupported custom forward paths fail explicitly instead of silently emitting an unoptimized graph.
This branch is updated with the latest
main.