Replace op.Constant(value_float=...) with Python float literals, use autocast

src/mobius/components/_ecapa_tdnn.py

@@ -313,7 +313,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
             op.ReduceSum(op.Mul(diff, diff), [2], keepdims=True),
             seq_length_float,
         )
-        eps_const = op.Constant(value_float=self._eps)
+        eps_const = self._eps
         std = op.Sqrt(op.Add(variance, eps_const))
 
         # Expand mean and std to (batch, channels, time) for concatenation

src/mobius/components/_multimodal.py

@@ -182,7 +182,7 @@ def forward(
             axis=0,
         )
         zero_pad = op.Expand(
-            op.CastLike(op.Constant(value_float=0.0), vision_embeddings),
+            op.CastLike(0.0, vision_embeddings),
             pad_shape,
         )
         # [batch, vision_seq + 1, hidden]

src/mobius/components/_qwen25_vl_vision.py

@@ -252,7 +252,7 @@ def _emit_standard_attention(self, op, q, k, v, cu_seqlens, seq_len_val):
         v = op.Unsqueeze(v, [0])
 
         # Build block-diagonal attention bias from cu_seqlens
-        attn_bias = self._build_block_diagonal_bias(op, cu_seqlens, seq_len_val)
+        attn_bias = self._build_block_diagonal_bias(op, cu_seqlens, seq_len_val, q)
         attn_bias = op.Unsqueeze(attn_bias, [0, 1])  # (1, 1, N, N)
 
         # Scaled dot-product attention
@@ -301,11 +301,14 @@ def _apply_rotary(self, op, x, cos, sin):
 
         return op.Concat(rot_x1, rot_x2, axis=-1)
 
-    def _build_block_diagonal_bias(self, op, cu_seqlens, seq_len):
+    def _build_block_diagonal_bias(self, op, cu_seqlens, seq_len, dtype_ref):
         """Build block-diagonal attention bias from cu_seqlens.
 
         Creates a matrix where positions in the same sub-sequence have 0
         and positions in different sub-sequences have -inf.
+
+        Args:
+            dtype_ref: reference tensor whose dtype the bias should match.
         """
         # Create range indices
         indices = op.Range(
@@ -337,8 +340,9 @@ def _build_block_diagonal_bias(self, op, cu_seqlens, seq_len):
         same_segment = op.Equal(seg_row, seg_col)
 
         # Convert to attention bias: 0 where same segment, -inf where different
-        neg_inf = op.Constant(value_float=-1e9)
-        zero = op.Constant(value_float=0.0)
+        # CastLike before Where so only the scalar is cast (cheaper than post-broadcast)
+        neg_inf = op.CastLike(-1e9, dtype_ref)
+        zero = op.CastLike(0.0, dtype_ref)
         return op.Where(same_segment, zero, neg_inf)
 
 

src/mobius/components/_qwen3_vl_vision.py

@@ -292,8 +292,8 @@ def _emit_standard_attention(self, op, query, key, value, cu_seqlens, hidden_sta
         same_segment = op.Equal(segment_row, segment_column)  # (total_seq, total_seq)
         attn_bias = op.Where(
             same_segment,
-            op.Constant(value_float=0.0),
-            op.Constant(value_float=-10000.0),
+            op.CastLike(0.0, query),
+            op.CastLike(-10000.0, query),
         )
         # Reshape for Attention: (1, 1, total_seq, total_seq)
         attn_bias = op.Unsqueeze(attn_bias, [0, 1])
@@ -609,12 +609,12 @@ def _interpolate_pos_embed(self, op, grid_thw):
         )
 
         h_idxs = body_op.Div(
-            body_op.Mul(h_range, body_op.Constant(value_float=n_minus_1)),
-            body_op.Sub(H_f, body_op.Constant(value_float=1.0)),
+            body_op.Mul(h_range, n_minus_1),
+            body_op.Sub(H_f, 1.0),
         )
         w_idxs = body_op.Div(
-            body_op.Mul(w_range, body_op.Constant(value_float=n_minus_1)),
-            body_op.Sub(W_f, body_op.Constant(value_float=1.0)),
+            body_op.Mul(w_range, n_minus_1),
+            body_op.Sub(W_f, 1.0),
         )
 
         # Floor/ceil indices
@@ -648,8 +648,8 @@ def _interpolate_pos_embed(self, op, grid_thw):
         idx_10 = body_op.Reshape(body_op.Add(bh_c2, wf2), [-1])
         idx_11 = body_op.Reshape(body_op.Add(bh_c2, wc2), [-1])
 
-        one_minus_dh = body_op.Sub(body_op.Constant(value_float=1.0), dh)
-        one_minus_dw = body_op.Sub(body_op.Constant(value_float=1.0), dw)
+        one_minus_dh = body_op.Sub(1.0, dh)
+        one_minus_dw = body_op.Sub(1.0, dw)
         dh2 = body_op.Unsqueeze(dh, [1])
         omdh2 = body_op.Unsqueeze(one_minus_dh, [1])
         dw2 = body_op.Unsqueeze(dw, [0])
@@ -704,7 +704,7 @@ def _interpolate_pos_embed(self, op, grid_thw):
             body_op.Constant(value_ints=[0]),
             axis=0,
         )
-        padded = body_op.Pad(pos_embeds, pads, body_op.Constant(value_float=0.0))
+        padded = body_op.Pad(pos_embeds, pads, 0.0)
         padded.name = "padded_pos_embed"
         body_graph.outputs.append(padded)
 

src/mobius/models/base.py

@@ -126,7 +126,7 @@ def forward(
                 to=ir.DataType.INT32,
             )  # [batch] INT32
             total_seq_len = op.Cast(
-                op.Gather(op.Shape(attention_mask), op.Constant(value_int=1)),
+                op.Gather(op.Shape(attention_mask), 1),
                 to=ir.DataType.INT32,
             )  # scalar INT32
 

src/mobius/models/blip2.py

@@ -166,7 +166,7 @@ def forward(self, op: builder.OpBuilder, input_ids: ir.Value, image_features: ir
 
         # Build indices to scatter image features into text positions
         mask_int = op.Cast(image_mask, to=7)
-        cumsum = op.CumSum(mask_int, op.Constant(value_int=1))
+        cumsum = op.CumSum(mask_int, 1)
         indices = op.Sub(cumsum, op.Constant(value_int=1))
         indices = op.Clip(indices, op.Constant(value_int=0))
 

src/mobius/models/clip.py

@@ -294,18 +294,20 @@ def forward(
         hidden_states = self.embeddings(op, input_ids)
 
         # Build causal attention bias (lower-triangular)
+        # CastLike ensures the bias dtype matches hidden_states (fp16/bf16/fp32)
+        # Cast the scalar *before* Expand so only a single element is cast
         seq_len = op.Shape(input_ids, start=1, end=2)
         _causal_mask = op.Trilu(
             op.Expand(
-                op.Constant(value_float=0.0),
+                op.CastLike(0.0, hidden_states),
                 op.Concat(seq_len, seq_len, axis=0),
             ),
             upper=0,
         )
         # Fill upper triangle with -inf
         neg_inf_mask = op.Trilu(
             op.Expand(
-                op.Constant(value_float=-10000.0),
+                op.CastLike(-10000.0, hidden_states),
                 op.Concat(seq_len, seq_len, axis=0),
             ),
             upper=1,

src/mobius/models/cogvideox.py

@@ -157,7 +157,7 @@ def forward(
             emb, num_outputs=6, axis=-1, _outputs=6
         )
 
-        one = op.Constant(value_float=1.0)
+        one = 1.0
 
         # Modulate video stream
         normed_h = self.norm(op, hidden_states)
@@ -197,7 +197,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, temb: ir.Value
         emb = self.linear(op, emb)
         # CogVideoX shift-first order
         shift, scale = op.Split(emb, num_outputs=2, axis=-1, _outputs=2)
-        one = op.Constant(value_float=1.0)
+        one = 1.0
         hidden_states = self.norm(op, hidden_states)
         hidden_states = op.Mul(hidden_states, op.Add(one, op.Unsqueeze(scale, [1])))
         hidden_states = op.Add(hidden_states, op.Unsqueeze(shift, [1]))

src/mobius/models/cohere.py

@@ -127,7 +127,5 @@ def forward(
             # Scale logits by the model's configured logit_scale scalar
             # (HF default: 0.0625 = 1/16 for all Cohere models).
             # CastLike ensures the constant matches logits dtype (fp16/bf16/fp32).
-            logits = op.Mul(
-                logits, op.CastLike(op.Constant(value_float=float(self.logit_scale)), logits)
-            )
+            logits = op.Mul(logits, op.CastLike(float(self.logit_scale), logits))
         return logits, present_key_values

src/mobius/models/deepseek.py

@@ -96,7 +96,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value):
         # Normalize weights (V3 with norm_topk_prob=True)
         if self.norm_topk_prob:
             weight_sum = op.ReduceSum(routing_weights, [-1], keepdims=True)
-            eps = op.Constant(value_float=1e-20)
+            eps = 1e-20
             routing_weights = op.Div(routing_weights, op.Add(weight_sum, eps))
 
         # Apply routing scale

src/mobius/models/deepseek_ocr2.py

@@ -396,13 +396,13 @@ def forward(
 
         # Cumulative sum to map flat image_features indices
         mask_int = op.Cast(image_mask, to=7)  # INT64
-        cumsum = op.CumSum(mask_int, op.Constant(value_int=1))
+        cumsum = op.CumSum(mask_int, 1)
         indices = op.Sub(cumsum, op.Constant(value_int=1))
         indices = op.Clip(indices, op.Constant(value_int=0))
 
         # Pad image_features for text-only safety
         pad_row = op.Expand(
-            op.Constant(value_float=0.0),
+            0.0,
             op.Concat(
                 op.Constant(value_ints=[1]),
                 op.Shape(image_features, start=1, end=2),

src/mobius/models/diffllama.py

@@ -175,7 +175,7 @@ def forward(
 
         # Scaled dot-product: Q_4d @ K_4d^T / sqrt(D) → [B, H, S, Sk]
         k_t = op.Transpose(k_4d, perm=[0, 1, 3, 2])  # [B, H, D, Sk]
-        scale = op.Constant(value_float=float(self._scaling))
+        scale = float(self._scaling)
         attn_scores = op.Mul(op.MatMul(q_4d, k_t), scale)
 
         # Causal mask: build [Sk, Sk] lower-triangle via Trilu, then slice
@@ -188,17 +188,18 @@ def forward(
         past_len = op.Sub(kv_len, q_len)  # [1]
         causal_slice = op.Slice(lower_tri, past_len, kv_len, [0])  # [S, Sk]
         causal_bool = op.Cast(causal_slice, to=9)  # BOOL (dtype 9)
-        zero_f = op.Constant(value_float=0.0)
-        neg_inf_f = op.Constant(value_float=float("-inf"))
+        # CastLike before Where so the scalar is cast while still 1-element
+        zero_f = op.CastLike(0.0, q_4d)
+        neg_inf_f = op.CastLike(float("-inf"), q_4d)
         causal_bias = op.Where(causal_bool, zero_f, neg_inf_f)  # [S, Sk]
         # Unsqueeze to [1, 1, S, Sk] and add to scores [B, H, S, Sk]
         attn_scores = op.Add(attn_scores, op.Unsqueeze(causal_bias, [0, 1]))
 
         # Padding mask (3D bool [B, S, Sk]) → additive bias [B, 1, S, Sk]
         if attention_bias is not None:
-            zero_f2 = op.Constant(value_float=0.0)
-            neg_inf_f2 = op.Constant(value_float=float("-inf"))
-            pad_bias = op.Where(attention_bias, zero_f2, neg_inf_f2)
+            pad_zero = op.CastLike(0.0, q_4d)
+            pad_neg_inf = op.CastLike(float("-inf"), q_4d)
+            pad_bias = op.Where(attention_bias, pad_zero, pad_neg_inf)
             attn_scores = op.Add(attn_scores, op.Unsqueeze(pad_bias, [1]))
 
         # Softmax and weighted sum with doubled V
@@ -213,7 +214,7 @@ def forward(
         #   lambda_full = exp(lq1·lk1) - exp(lq2·lk2) + lambda_init
         lam1 = op.Exp(op.ReduceSum(op.Mul(self.lambda_q1, self.lambda_k1), keepdims=False))
         lam2 = op.Exp(op.ReduceSum(op.Mul(self.lambda_q2, self.lambda_k2), keepdims=False))
-        lam_init = op.Constant(value_float=float(self._lambda_init))
+        lam_init = float(self._lambda_init)
         lam_full = op.Add(op.Sub(lam1, lam2), lam_init)
 
         # Differential output: out1 - lambda_full * out2  →  [B, H/2, S, 2D]
@@ -222,9 +223,9 @@ def forward(
         # GroupNorm = RMSNorm without learnable affine params (on last dim = 2D)
         #   rms_norm(x) = x / sqrt(mean(x²) + eps)
         x_sq = op.Mul(diff, diff)
-        eps = op.Constant(value_float=float(self._rms_norm_eps))
+        eps = float(self._rms_norm_eps)
         rms = op.Sqrt(op.Add(op.ReduceMean(x_sq, [-1], keepdims=True), eps))
-        scale_norm = op.Constant(value_float=float(1.0 - self._lambda_init))
+        scale_norm = float(1.0 - self._lambda_init)
         normed = op.Mul(op.Div(diff, rms), scale_norm)
 
         # Reshape [B, H/2, S, 2D] → [B, S, H*D]

src/mobius/models/dit.py

@@ -148,7 +148,7 @@ def forward(
         )
 
         # Apply scale and shift to normed hidden_states for self-attention
-        one = op.Constant(value_float=1.0)
+        one = 1.0
         attn_input = op.Mul(normed, op.Add(one, op.Unsqueeze(scale_msa, [1])))
         attn_input = op.Add(attn_input, op.Unsqueeze(shift_msa, [1]))
 

src/mobius/models/falcon.py

@@ -293,10 +293,11 @@ def _create_alibi_bias(op, num_heads: int, seq_len, total_len):
     alibi = op.Mul(slopes_4d, bias_2d)  # [1, num_heads, seq_len, total_len]
 
     # Causal mask: mask future positions with large negative value
+    # CastLike before Where so only the scalar is cast (cheaper than post-broadcast)
     causal_mask = op.Where(
         op.GreaterOrEqual(q_with_offset, kv_expanded),
-        op.Constant(value_float=0.0),
-        op.Constant(value_float=-10000.0),
+        op.CastLike(0.0, neg_distance),
+        op.CastLike(-10000.0, neg_distance),
     )  # [seq_len, total_len]
     causal_4d = op.Unsqueeze(causal_mask, [0, 1])  # [1, 1, seq_len, total_len]
     return op.Add(alibi, causal_4d)

src/mobius/models/flux_sd3.py

@@ -119,7 +119,7 @@ def forward(
             hidden_states,
             temb,
         )
-        one = op.Constant(value_float=1.0)
+        one = 1.0
         img_input = op.Mul(normed, op.Add(one, op.Unsqueeze(scale_msa, [1])))
         img_input = op.Add(img_input, op.Unsqueeze(shift_msa, [1]))
 
@@ -324,7 +324,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, temb: ir.Value
             hidden_states,
             temb,
         )
-        one = op.Constant(value_float=1.0)
+        one = 1.0
         attn_input = op.Mul(normed, op.Add(one, op.Unsqueeze(scale_msa, [1])))
         attn_input = op.Add(attn_input, op.Unsqueeze(shift_msa, [1]))
 

src/mobius/models/gemma3.py

@@ -135,7 +135,7 @@ def forward(self, op: builder.OpBuilder, input_ids: ir.Value, image_features: ir
         image_mask_3d = op.Unsqueeze(image_mask, [-1])
 
         mask_int = op.Cast(image_mask, to=7)
-        cumsum = op.CumSum(mask_int, op.Constant(value_int=1))
+        cumsum = op.CumSum(mask_int, 1)
         indices = op.Sub(cumsum, op.Constant(value_int=1))
         indices = op.Clip(indices, op.Constant(value_int=0))
 

src/mobius/models/hunyuan_dit.py

@@ -109,7 +109,7 @@ def forward(self, op: builder.OpBuilder, x: ir.Value):
         mean = op.ReduceMean(x, [-1], keepdims=True)
         diff = op.Sub(x, mean)
         var = op.ReduceMean(op.Mul(diff, diff), [-1], keepdims=True)
-        eps = op.Constant(value_float=self._eps)
+        eps = self._eps
         return op.Div(diff, op.Sqrt(op.Add(var, eps)))
 
 
@@ -383,7 +383,7 @@ def forward(self, op: builder.OpBuilder, hidden_states: ir.Value, temb: ir.Value
         scale, shift = op.Split(emb, num_outputs=2, axis=-1, _outputs=2)
         # Norm → (1 + scale) * normed + shift
         normed = self.norm(op, hidden_states)
-        one = op.Constant(value_float=1.0)
+        one = 1.0
         normed = op.Mul(normed, op.Add(one, op.Unsqueeze(scale, [1])))
         return op.Add(normed, op.Unsqueeze(shift, [1]))
 

src/mobius/models/internvl.py

@@ -434,11 +434,11 @@ def _pixel_shuffle(self, op, x):
 
         # Step 1: view(N, W, H*scale, C/scale)
         h_scaled = op.Cast(
-            op.Mul(op.Cast(h, to=1), op.Constant(value_float=scale)),
+            op.Mul(op.Cast(h, to=1), scale),
             to=7,
         )
         c_over_scale = op.Cast(
-            op.Div(op.Cast(channels, to=1), op.Constant(value_float=scale)),
+            op.Div(op.Cast(channels, to=1), scale),
             to=7,
         )
         shape_step1 = op.Concat(batch, w, h_scaled, c_over_scale, axis=0)
@@ -449,13 +449,13 @@ def _pixel_shuffle(self, op, x):
 
         # Step 3: view(N, H*scale, W*scale, C/(scale^2))
         w_scaled = op.Cast(
-            op.Mul(op.Cast(w, to=1), op.Constant(value_float=scale)),
+            op.Mul(op.Cast(w, to=1), scale),
             to=7,
         )
         c_over_scale2 = op.Cast(
             op.Div(
                 op.Cast(channels, to=1),
-                op.Constant(value_float=scale * scale),
+                scale * scale,
             ),
             to=7,
         )
@@ -521,7 +521,7 @@ def forward(
 
         # Compute indices into image_features via cumulative sum
         mask_int = op.Cast(image_mask, to=7)
-        cumsum = op.CumSum(mask_int, op.Constant(value_int=1))
+        cumsum = op.CumSum(mask_int, 1)
         indices = op.Sub(cumsum, op.Constant(value_int=1))
         indices = op.Clip(indices, op.Constant(value_int=0))
 

src/mobius/models/llava.py

@@ -125,15 +125,15 @@ def forward(self, op: builder.OpBuilder, input_ids: ir.Value, image_features: ir
         image_mask_3d = op.Unsqueeze(image_mask, [-1])
 
         mask_int = op.Cast(image_mask, to=7)
-        cumsum = op.CumSum(mask_int, op.Constant(value_int=1))
+        cumsum = op.CumSum(mask_int, 1)
         indices = op.Sub(cumsum, op.Constant(value_int=1))
         indices = op.Clip(indices, op.Constant(value_int=0))
 
         # Pad image_features with one zero row so Gather is valid even when
         # image_features is empty (text-only input: num_image_tokens == 0).
         # The Where mask ensures the padding row is never used in the output.
         pad_row = op.Expand(
-            op.CastLike(op.Constant(value_float=0.0), image_features),
+            op.CastLike(0.0, image_features),
             op.Concat(
                 op.Constant(value_ints=[1]),
                 op.Shape(image_features, start=1, end=2),

src/mobius/models/qwen3_asr.py

@@ -252,7 +252,7 @@ def forward(
         is_audio_int = op.Cast(is_audio, to=7)  # INT64
         # Flatten across batch for cumsum then reshape
         flat_mask = op.Reshape(is_audio_int, op.Constant(value_ints=[-1]))
-        flat_indices = op.CumSum(flat_mask, op.Constant(value_int=0))
+        flat_indices = op.CumSum(flat_mask, 0)
         flat_indices = op.Mul(flat_indices, flat_mask)
         # Reshape back to (batch, seq_len)
         indices = op.Reshape(flat_indices, op.Shape(input_ids))