Replace op.Constant(value_float=...) with Python float literals, use autocast - #166
Merged
Conversation
…rals Replace all 68 instances of op.Constant(value_float=X) across 30 files with plain Python float literals or float() expressions. This leverages onnxscript's auto-casting of Python scalars to match the dtype of the other operand in binary ops, which fixes dtype mismatches when models use bfloat16 or float16. Patterns replaced: - op.Constant(value_float=X) -> X (plain literal) - op.Constant(value_float=float(expr)) -> float(expr) - op.Constant(value_float=self.attr) -> self.attr - op.CastLike(op.Constant(value_float=X), ref) -> op.CastLike(X, ref) Note: op.Constant(value_int=X) replacements were NOT included because onnxscript creates initializers (not Constant nodes) for Python int literals, causing 'already registered' collisions when the same int value appears multiple times in a graph. This is an onnxscript limitation that needs to be resolved upstream before value_int cleanup can proceed. Continuation of PR #58 which established this pattern. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
Performance Comparison
|
|
The author of this PR, gramalingam, is not an activated member of this organization on Codecov. |
🏗️ Architecture Diff
qwen3_5_vl (hybrid-qwen-vl) / embedding — 8 change(s)Op summary: 20 → 19 nodes --- base
+++ head
@@ -9,7 +9,6 @@
Sub
Constant
Clip
-Constant
CastLike
Constant
ShapeRemoved nodes:
Connectivity changes:
Initializer changes:
qwen3_5_vl (hybrid-qwen-vl) / vision — 212 change(s)Op summary: 240 → 240 nodes --- base
+++ head
@@ -134,8 +134,8 @@
Unsqueeze
Unsqueeze
Equal
-Constant
-Constant
+CastLike
+CastLike
Where
Unsqueeze
Unsqueeze
@@ -207,8 +207,8 @@
Unsqueeze
Unsqueeze
Equal
-Constant
-Constant
+CastLike
+CastLike
Where
Unsqueeze
UnsqueezeAdded nodes:
Removed nodes:
Connectivity changes:
Initializer changes:
t5 (seq2seq) / decoder — 39 change(s)Op summary: 99 → 95 nodes --- base
+++ head
@@ -19,13 +19,9 @@
Constant
Less
Cast
-Constant
Max
-Constant
Div
Log
-Constant
-Constant
Mul
Add
CastRemoved nodes:
Modified attributes:
Connectivity changes:
Initializer changes:
t5 (seq2seq) / encoder — 28 change(s)Op summary: 75 → 71 nodes --- base
+++ head
@@ -19,13 +19,9 @@
Constant
Less
Cast
-Constant
Max
-Constant
Div
Log
-Constant
-Constant
Mul
Add
CastRemoved nodes:
Modified attributes:
Connectivity changes:
Initializer changes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR continues the scalar-constant refactor across the ONNX graph builders by replacing many op.Constant(value_float=...) usages with Python float literals (or float(...)) to rely on onnxscript’s scalar auto-casting for fp16/bf16 friendliness.
Changes:
- Replaced many
op.Constant(value_float=...)float scalars with Python float literals /float(...)expressions across models, components, and functions. - Simplified several
op.CastLike(op.Constant(...), ref)patterns toop.CastLike(<python-float>, ref). - Also includes some
value_int→ Python int literal substitutions in a few places.
Reviewed changes
Copilot reviewed 35 out of 35 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/models/vae.py | Replace attention scale float Constant with Python float. |
| src/mobius/models/t5.py | Replace float Constants in relative-position bucketing and output scaling. |
| src/mobius/models/segformer.py | Replace attention scale float Constant with attribute/literal. |
| src/mobius/models/sam2.py | Replace Gather indices from op.Constant(value_int=...) to Python ints. |
| src/mobius/models/qwen_vl.py | Replace scalar Constants (floats/ints) in embedding + DeepStack injection paths. |
| src/mobius/models/qwen_image_vae.py | Replace Pad/Clip float Constants with float literals. |
| src/mobius/models/qwen_image.py | Replace float Constants (e.g., 1.0, timestep scaling) with literals. |
| src/mobius/models/qwen3_tts_tokenizer.py | Replace Clip bounds / scalar multiplier float Constant with literals. |
| src/mobius/models/qwen3_asr.py | Replace CumSum axis value_int Constant with Python int. |
| src/mobius/models/minimax.py | Simplify CastLike of alpha/beta float constants. |
| src/mobius/models/llava.py | Replace CumSum axis constant and zero-pad CastLike constant pattern. |
| src/mobius/models/internvl.py | Replace pixel-shuffle scale float Constants and CumSum axis constant. |
| src/mobius/models/hunyuan_dit.py | Replace eps/one float Constants with literals. |
| src/mobius/models/granitemoehybrid.py | Simplify CastLike residual multiplier constants. |
| src/mobius/models/gemma3.py | Replace CumSum axis constant with Python int. |
| src/mobius/models/flux_sd3.py | Replace one float Constants with literals. |
| src/mobius/models/falcon.py | Replace Where branch float Constants with literals in ALiBi bias. |
| src/mobius/models/dit.py | Replace one float Constant with literal. |
| src/mobius/models/diffllama.py | Replace scale/eps/lambda float Constants with literals; adjust masking constants. |
| src/mobius/models/deepseek_ocr2.py | Replace CumSum axis constant and zero padding Constant with literal. |
| src/mobius/models/deepseek.py | Replace small epsilon float Constant with literal. |
| src/mobius/models/cohere.py | Simplify CastLike around logit_scale float Constant. |
| src/mobius/models/cogvideox.py | Replace one float Constants with literals. |
| src/mobius/models/clip.py | Replace Expand float Constants for causal masks with literals. |
| src/mobius/models/blip2.py | Replace CumSum axis constant with Python int. |
| src/mobius/models/base.py | Replace Gather index value_int Constant with Python int. |
| src/mobius/functions/packed_multi_head_attention.py | Replace Where branch float Constants with literals in attention bias construction. |
| src/mobius/functions/linear_attention.py | Simplify CastLike(scale) constant usage. |
| src/mobius/functions/causal_conv.py | Replace Gather index and Sub constant input from value_int to Python ints. |
| src/mobius/components/_qwen3_vl_vision.py | Replace Where/Pad float Constants with literals in attention bias/interpolation paths. |
| src/mobius/components/_qwen25_vl_vision.py | Replace bias scalars (neg_inf, zero) from Constants to literals. |
| src/mobius/components/_multimodal.py | Simplify CastLike(0.0) constant usage. |
| src/mobius/components/_mamba_block.py | Simplify CastLike(0.0) and Clip min constant usage. |
| src/mobius/components/_gated_deltanet.py | Simplify CastLike(0.0) constant usage. |
| src/mobius/components/_ecapa_tdnn.py | Replace eps float Constant with literal. |
Comments suppressed due to low confidence (2)
src/mobius/models/deepseek_ocr2.py:412
pad_rowis created viaop.Expand(0.0, ...)and then concatenated withimage_features. Ifimage_featuresis fp16/bf16, the scalar0.0typically materializes as a float32 tensor (Expand has no typed float tensor operand), andop.Concatwill error due to mismatched element types. Useop.CastLike(0.0, image_features)(or castpad_rowtoimage_featuresdtype) before concatenation.
# Pad image_features for text-only safety
pad_row = op.Expand(
0.0,
op.Concat(
op.Constant(value_ints=[1]),
op.Shape(image_features, start=1, end=2),
axis=0,
),
)
padded_features = op.Concat(image_features, pad_row, axis=0)
src/mobius/components/_qwen3_vl_vision.py:300
attn_biasis produced byop.Where(..., 0.0, -10000.0)with only a boolean tensor as typed input, so the bias tensor will typically default to float32.op.Attentionexpects the bias dtype to matchquery/key/value(often fp16/bf16 in this model), so this can create a type mismatch. Castattn_biastoquery's dtype (e.g.,op.CastLike(attn_bias, query)) before the Attention call.
attn_bias = op.Where(
same_segment,
0.0,
-10000.0,
)
# Reshape for Attention: (1, 1, total_seq, total_seq)
attn_bias = op.Unsqueeze(attn_bias, [0, 1])
Address PR #166 review feedback: 1. Add op.CastLike() on float scalars BEFORE Where/Expand ops so the attention bias dtype matches the model's compute dtype (fp16/bf16). Cast is applied to the scalar (1-element) before broadcast, which is cheaper than casting after expansion. Affected files: - clip.py: CastLike 0.0/-10000.0 before Expand in causal bias - diffllama.py: CastLike 0.0/-inf before Where in causal + pad bias - falcon.py: CastLike 0.0/-10000.0 before Where in ALiBi causal mask - packed_multi_head_attention.py: CastLike before Where in segment bias - _qwen25_vl_vision.py: CastLike before Where in block-diagonal bias - _qwen3_vl_vision.py: CastLike before Where in block-diagonal bias 2. Revert value_int literal changes (out of PR scope, risk of initializer name collisions due to onnxscript limitation): - sam2.py: revert Gather indices back to op.Constant(value_int=N) - qwen_vl.py: revert CumSum axis back to op.Constant(value_int=1) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
justinchuby
approved these changes
Apr 14, 2026
Python float/int literals passed to onnxscript ops become initializers in the root graph. This breaks in two contexts: 1. ir.Function bodies: functions cannot reference outer-scope initializers (unlike subgraphs which can). Affected: - linear_attention.py (scale in LinearAttention function) - packed_multi_head_attention.py (bias in function body) - causal_conv.py (axis/width in CausalConvWithState function) 2. Component code that may run in subgraph or function contexts: - minimax.py (_scaled_add in MoE expert dispatch) - granitemoehybrid.py (residual multiplier in MoE+Mamba layers) - _mamba_block.py (conv bias, time_step_min in Scan body) - _gated_deltanet.py (conv bias) Reverts these files to use op.Constant(value_float=...) which creates inline Constant nodes that live in whatever graph context they are built in. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam <grama@microsoft.com>
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.
Continuation of PR #58 which established this pattern.
Replace all 68 instances of op.Constant(value_float=X) across 30 files with plain Python float literals or float() expressions. This leverages onnxscript's auto-casting of Python scalars to match the dtype of the other operand in binary ops, which fixes dtype mismatches when models use bfloat16 or float16.
Patterns replaced:
Note: This is part 1. (A similar change can be done for int/ints/floats, but will do that in separate PRs.)