fix: restore Shape(input_ids) for query_length in create_attention_bias()

Commit fd85bdf changed query_length to use Shape(attention_mask, start=1, end=2)
so that the EliminateShape WebGPU rule would fire on it.  But attention_mask
always has shape (batch, total_seq_len), so:

  query_length = Shape(attention_mask, 1)  →  total_seq_len
  total_length = Shape(attention_mask, 1)  →  total_seq_len
  start = Sub(total_length, query_length)  →  0

start=0 means Slice selects ALL positions as queries.  During decode,
input_ids is (batch, 1), so query_length should be 1 and start should be
total_seq_len - 1, slicing only the last row.  With start=0 the returned
bias has shape (batch, 1, total_seq_len, total_seq_len) instead of the
correct (batch, 1, 1, total_seq_len), causing completely wrong attention
scores and garbage model output.

Fix: revert query_length to Shape(input_ids, start=1, end=2).
- total_length stays on attention_mask (EliminateShape still fires for it).
- On WebGPU (concrete dims), Shape(input_ids, 1) is constant-folded away,
  so no dynamic Shape op survives in the WebGPU export.

Add test: test_query_length_from_input_ids_not_attention_mask verifies
that a Shape node reading from input_ids is present in the graph, guarding
against future regressions.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

src/mobius/components/_common.py

@@ -191,10 +191,14 @@ def create_attention_bias(
     # Actually we need to implement this with shape ops
 
     # Get query_length and total_length from shapes.
-    # Both come from attention_mask: total_length is the full sequence length, and
-    # query_length is derived as total_length - past_length.  Using attention_mask
-    # for both means the EliminateShape rule on WebGPU can eliminate these Shape ops.
-    query_length = op.Shape(attention_mask, start=1, end=2)  # 1-D [1]
+    # query_length comes from input_ids dim 1 (the query; e.g. 1 during decode).
+    # total_length comes from attention_mask dim 1 (past + current tokens).
+    # Using attention_mask for total_length lets the EliminateShape WebGPU rule
+    # eliminate that Shape op.  Using input_ids for query_length is semantically
+    # correct: during decode input_ids is (batch, 1), so query_length=1 and
+    # start = total_length - 1, giving the last row of q_indices.
+    # On WebGPU (concrete dims), Shape(input_ids, 1) is constant-folded away.
+    query_length = op.Shape(input_ids, start=1, end=2)  # 1-D [1]
     total_length = op.Shape(attention_mask, start=1, end=2)  # 1-D [1]
     start = op.Sub(total_length, query_length)
     # q_indices_2d: (batch_size, query_length)

src/mobius/components/_common_test.py

@@ -91,6 +91,31 @@ def test_creates_bias_with_sliding_window(self):
         assert bias is not None
         assert count_op_type(graph, "Less") >= 1
 
+    def test_query_length_from_input_ids_not_attention_mask(self):
+        """Shape(input_ids, 1) must provide query_length, not Shape(attention_mask, 1).
+
+        During decode, input_ids is (batch, 1) and attention_mask is (batch, total_len).
+        Both shapes must produce separate Shape nodes so the Slice picks only the
+        last query row.  If both came from attention_mask, start=0 and the full
+        sequence would be used as queries, producing wrong attention scores.
+        """
+        builder, op, graph = create_test_builder()
+        # Simulate decode: q_len=1, total_len=8 (7 past + 1 current token)
+        input_ids = create_test_input(builder, "input_ids", [2, 1], dtype=ir.DataType.INT64)
+        attention_mask = create_test_input(
+            builder, "attention_mask", [2, 8], dtype=ir.DataType.INT64
+        )
+        create_attention_bias(op, input_ids, attention_mask)
+
+        # query_length must come from input_ids (dim 1 = 1), not attention_mask.
+        # Verify there is a Shape node that reads from input_ids.
+        shape_inputs = [
+            n.inputs[0].name for n in graph if n.op_type == "Shape" and n.inputs[0] is not None
+        ]
+        assert any(name == "input_ids" for name in shape_inputs), (
+            "Shape(input_ids, 1) must be present to extract query_length correctly"
+        )
+
 
 class TestCreatePaddingMask:
     def test_creates_bool_mask_with_2d_input_ids(self):