Fix Qwen2.5/3-VL vision PackedMultiHeadAttention crash on CUDA (#358)

## Problem

Building **Qwen2.5-VL-3B** with Mobius and evaluating on **CUDA**
crashes at the first windowed vision block:

```
[E ... PackedMultiHeadAttention node 'vision_encoder/visual/blocks.0/attn/PackedMultiHeadAttention_node_214']
Status Message: Input 'cumulative_sequence_length' should have 1 dimension with size equal to batch_size + 1
```

## Root cause

The EP-gated packed path (`_emit_packed_mha`, used when the EP
advertises `supports_packed_multi_head_attention`, i.e. CUDA / trt-rtx)
built `token_offset` as shape **`(1, N)`**, declaring `batch_size = 1`
to `com.microsoft::PackedMultiHeadAttention`.

ORT derives `batch_size` from `token_offset.shape[0]` and requires
`cumulative_sequence_length` (our `cu_seqlens`) to have length
`batch_size + 1`
(`contrib_ops/cuda/bert/packed_multihead_attention.cc`). Whenever
`cu_seqlens` enumerates more than one sub-sequence:

- **windowed** vision blocks (`cu_window_seqlens`, length `num_windows +
1`), or
- multi-frame **video** / multi-image **full-attention** blocks (length
`T + 1`),

`num_sub_seqs + 1 > 2`, so the kernel rejects the input and the run
crashes. The Copilot diagnosis in the issue was accurate.

Even in the non-crashing single-sub-sequence case, `batch_size = 1` made
the kernel compute **full** attention over all `N` tokens instead of the
intended **block-diagonal** attention — so the path was doubly broken.

`_qwen3_vl_vision.py` had the identical copy-pasted bug (latent crash
for multi-image / video).

## Fix

Build `token_offset` with shape `(num_sub_seqs, max_seq_len)` following
ORT's `GetPaddingOffset` convention — valid padded-layout token indices
(`b * max_seq_len + s`) in packed order, then padding-slot indices. This
matches the already-correct `BlockDiagonalToPackedMHA` rewrite rule,
extracted into a shared `build_packed_token_offset` helper in
`components/_common.py` and reused by both vision encoders.

Also stop unconditionally down-casting q/k/v to float16: the PackedMHA
CUDA kernel supports **both float32 and float16**, so only **bfloat16**
builds (unsupported) are cast to f16; f32/f16 stay native to preserve
precision.

## Verification

- **CUDA parity:** with the fix, the packed path matches the standard
block-diagonal attention path (CPU) on a 4-window grid
(`grid_thw=[1,12,12]`) that crashed before — **cosine 1.0, max abs diff
~7e-8**.
- New exact-value unit tests for `build_packed_token_offset` against
ORT's own reference data (e.g. `cu=[0,1,3] -> [[0,2],[3,1]]`, including
padding indices `>= token_count`, single-subsequence identity, uniform
windows).
- Full non-integration suite green (1318 passed); existing `qwen2_5_vl`
/ `qwen3_vl` graph tests and packed-attention rewrite tests unaffected.

---------

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

Author: justinchuby

src/mobius/components/_common.py

@@ -264,6 +264,79 @@ def create_attention_bias(
     return op.Unsqueeze(attention_bias, [1])
 
 
+def build_packed_token_offset(op: OpBuilder, cu_seqlens) -> ir.Value:
+    """Build ``token_offset`` for ``com.microsoft::PackedMultiHeadAttention``.
+
+    ORT's ``PackedMultiHeadAttention`` treats the packed
+    ``(token_count, hidden)`` query/key/value as ``batch_size``
+    variable-length sub-sequences (delimited by ``cu_seqlens``) with
+    right-padding removed, and computes block-diagonal attention *within*
+    each sub-sequence.  The kernel derives ``batch_size`` from
+    ``token_offset.shape[0]`` and requires ``cumulative_sequence_length``
+    (i.e. ``cu_seqlens``) to have length ``batch_size + 1``.
+
+    ``token_offset`` has shape ``(batch_size, max_seq_len)`` and follows
+    ORT's ``GetPaddingOffset`` convention: the first ``token_count`` entries
+    are the padded-layout indices (``b * max_seq_len + s``) of the valid
+    tokens in packed (sub-sequence-major) order, followed by the
+    padded-layout indices of the padding slots.
+
+    Example: ``cu_seqlens = [0, 1, 3]`` (two sub-sequences of length 1 and
+    2, so ``max_seq_len = 2``) yields ``[[0, 2], [3, 1]]``.
+
+    Args:
+        op: The OpBuilder.
+        cu_seqlens: ``(batch_size + 1,)`` cumulative sequence lengths
+            (INT32 or INT64).
+
+    Returns:
+        ``(batch_size, max_seq_len)`` INT32 ``token_offset`` tensor.
+    """
+    # Index/axis tensors are materialised as explicit Constant nodes (rather
+    # than relying on Python-list auto-conversion) so this helper works under
+    # both the component OpBuilder and the onnxscript rewriter op context.
+    axes_0 = op.Constant(value_ints=[0])
+    axes_1 = op.Constant(value_ints=[1])
+    neg_one = op.Constant(value_ints=[-1])
+    start_1 = op.Constant(value_ints=[1])
+    int_max = op.Constant(value_ints=[INT64_MAX])
+
+    cu_seqlens_i32 = op.Cast(cu_seqlens, to=ir.DataType.INT32)
+
+    # batch_size = len(cu_seqlens) - 1  (number of packed sub-sequences).
+    batch_size = op.Cast(
+        op.Sub(op.Size(cu_seqlens), op.Constant(value_int=1)),
+        to=ir.DataType.INT32,
+    )
+
+    # Per-sub-sequence lengths and the padded sequence dimension.
+    starts = op.Slice(cu_seqlens_i32, axes_0, neg_one, axes_0)  # cu[:-1]
+    ends = op.Slice(cu_seqlens_i32, start_1, int_max, axes_0)  # cu[1:]
+    lengths = op.Sub(ends, starts)  # (batch_size,)
+    max_len = op.Squeeze(op.ReduceMax(lengths), axes_0)  # scalar INT32
+
+    # Padded position grid: pos[b, s] = b * max_len + s.
+    zero_i32 = op.Cast(op.Constant(value_int=0), to=ir.DataType.INT32)
+    one_i32 = op.Cast(op.Constant(value_int=1), to=ir.DataType.INT32)
+    rows = op.Range(zero_i32, batch_size, one_i32)  # (batch_size,)
+    cols = op.Range(zero_i32, max_len, one_i32)  # (max_len,)
+    pos_matrix = op.Add(
+        op.Mul(op.Unsqueeze(rows, axes_1), max_len),
+        op.Unsqueeze(cols, axes_0),
+    )  # (batch_size, max_len) INT32
+    pos_matrix_shape = op.Shape(pos_matrix)
+
+    # Column s is a valid token for row b iff s < lengths[b].
+    valid_mask = op.Less(op.Unsqueeze(cols, axes_0), op.Unsqueeze(lengths, axes_1))
+    valid_mask_1d = op.Reshape(valid_mask, neg_one)
+    pos_1d = op.Reshape(pos_matrix, neg_one)
+
+    # Valid positions first (packed order), then padding-slot positions.
+    valid_indices = op.Compress(pos_1d, valid_mask_1d)
+    padding_indices = op.Compress(pos_1d, op.Not(valid_mask_1d))
+    return op.Reshape(op.Concat(valid_indices, padding_indices, axis=0), pos_matrix_shape)
+
+
 def create_padding_mask(
     op: OpBuilder,
     input_ids,

src/mobius/components/_common_test.py

@@ -13,6 +13,7 @@
 from mobius.components._common import (
     Embedding,
     Linear,
+    build_packed_token_offset,
     create_attention_bias,
     create_padding_mask,
     create_sliding_window_mask,
@@ -214,6 +215,65 @@ def test_decode_single_query_is_causal(self):
         assert bool((out[0, 0, 0] > -1.0).all())
 
 
+class TestBuildPackedTokenOffset:
+    """``build_packed_token_offset`` must reproduce ORT's GetPaddingOffset.
+
+    ORT's ``PackedMultiHeadAttention`` derives ``batch_size`` from
+    ``token_offset.shape[0]`` and requires ``cumulative_sequence_length`` to
+    have length ``batch_size + 1``.  ``token_offset`` lists the padded-layout
+    indices (``b * max_len + s``) of valid tokens (packed order) first, then
+    the padding slots.  We run the helper through ORT and compare exact values.
+    """
+
+    @staticmethod
+    def _build(dtype=ir.DataType.INT64):
+        b, op, g = create_test_builder()
+        cu = create_test_input(b, "cu_seqlens", ["K"], dtype=dtype)
+        token_offset = build_packed_token_offset(op, cu)
+        token_offset.name = "token_offset"
+        g.outputs.append(token_offset)
+        return ir.Model(g, ir_version=10)
+
+    def _run(self, cu_seqlens, np_dtype=np.int64, ir_dtype=ir.DataType.INT64):
+        sess = OnnxModelSession(self._build(ir_dtype), device="cpu")
+        return sess.run({"cu_seqlens": np.array(cu_seqlens, dtype=np_dtype)})["token_offset"]
+
+    def test_ort_reference_example(self):
+        # ORT test data: cu=[0,1,3] (lengths 1,2; max_len=2) -> [[0,2],[3,1]].
+        out = self._run([0, 1, 3])
+        assert out.dtype == np.int32
+        assert np.array_equal(out, np.array([[0, 2], [3, 1]], dtype=np.int32))
+
+    def test_padding_indices_exceed_token_count(self):
+        # cu=[0,2,5]: lengths [2,3], max_len=3, token_count=5.
+        # Padded grid pos = b*3 + s -> row0 valid cols {0,1} pad col {2};
+        # row1 valid cols {3,4,5}. valid (packed order) = [0,1,3,4,5];
+        # padding slot = [2]. token_offset = [[0,1,3],[4,5,2]].
+        out = self._run([0, 2, 5])
+        assert np.array_equal(out, np.array([[0, 1, 3], [4, 5, 2]], dtype=np.int32))
+        # Padding value (2) is a padded-layout index, here < token_count, but
+        # the construction may yield values >= token_count for other shapes.
+
+    def test_single_subsequence_is_identity(self):
+        # cu=[0,4]: one sub-sequence -> shape (1,4), identity [0,1,2,3].
+        out = self._run([0, 4])
+        assert out.shape == (1, 4)
+        assert np.array_equal(out, np.array([[0, 1, 2, 3]], dtype=np.int32))
+
+    def test_uniform_windows(self):
+        # Three windows of equal length 2: max_len=2, no padding.
+        out = self._run([0, 2, 4, 6])
+        assert out.shape == (3, 2)
+        assert np.array_equal(out, np.array([[0, 1], [2, 3], [4, 5]], dtype=np.int32))
+
+    def test_int32_input(self):
+        # The helper documents INT32 or INT64 cu_seqlens; INT32 input must
+        # produce the same result as the INT64 reference example.
+        out = self._run([0, 1, 3], np_dtype=np.int32, ir_dtype=ir.DataType.INT32)
+        assert out.dtype == np.int32
+        assert np.array_equal(out, np.array([[0, 2], [3, 1]], dtype=np.int32))
+
+
 class TestCreatePaddingMask:
     def test_creates_bool_mask_with_2d_input_ids(self):
         """Standard path: input_ids is 2D [batch, q_len]."""

src/mobius/components/_qwen25_vl_vision.py

@@ -26,8 +26,8 @@
 import onnx_ir as ir
 from onnxscript import OpBuilder, nn
 
-from mobius._build_context import ep_capabilities
-from mobius.components._common import LayerNorm, Linear
+from mobius._build_context import ep_capabilities, get_build_dtype
+from mobius.components._common import LayerNorm, Linear, build_packed_token_offset
 from mobius.components._mlp import FCMLP, GatedMLP
 from mobius.components._rms_norm import RMSNorm
 from mobius.components._scan_utils import (
@@ -189,10 +189,13 @@ def _emit_packed_mha(self, op, q, k, v, cu_seqlens, seq_len_val):
         """Emit com.microsoft.PackedMultiHeadAttention.
 
         Uses cu_seqlens natively, avoiding the O(N^2) block-diagonal bias.
+        The block-diagonal masking is expressed through the packed varlen
+        contract: each sub-sequence delimited by ``cu_seqlens`` (a window or
+        a frame) is one batch element, and attention is computed within it.
 
         Args:
             q, k, v: (N, num_heads, head_dim) after rotary embedding
-            cu_seqlens: (num_sub_seqs + 1,) INT32
+            cu_seqlens: (num_sub_seqs + 1,) cumulative sequence lengths
             seq_len_val: (1,) shape tensor with N
         """
         hidden_size = self.num_heads * self.head_dim
@@ -201,25 +204,25 @@ def _emit_packed_mha(self, op, q, k, v, cu_seqlens, seq_len_val):
         key = op.Reshape(k, op.Concat(seq_len_val, [hidden_size], axis=0))
         value = op.Reshape(v, op.Concat(seq_len_val, [hidden_size], axis=0))
 
-        # token_offset: identity mapping for packed (no-padding) input.
-        # Shape: (1, token_count) — single batch, positions [0..N-1].
-        token_count_scalar = op.Squeeze(seq_len_val, [0])
-        token_offset = op.Unsqueeze(
-            op.Range(
-                op.Constant(value_int=0),
-                token_count_scalar,
-                op.Constant(value_int=1),
-            ),
-            [0],
-        )
-        token_offset = op.Cast(token_offset, to=6)  # INT32
+        # token_offset: (num_sub_seqs, max_seq_len) mapping the packed tokens to
+        # their padded (batch, seq) layout. ORT derives batch_size from
+        # token_offset.shape[0] and requires cumulative_sequence_length to have
+        # length batch_size + 1, so this MUST encode every sub-sequence (window
+        # or frame), not a single (1, N) batch — otherwise the kernel rejects
+        # windowed cu_seqlens and computes full instead of block-diagonal
+        # attention.
+        token_offset = build_packed_token_offset(op, cu_seqlens)
 
         cu_seqlens_int32 = op.Cast(cu_seqlens, to=6)  # INT32
 
-        # PackedMHA doesn't support bfloat16; cast to float16 if needed
-        query_mha = op.Cast(query, to=10)  # FLOAT16
-        key_mha = op.Cast(key, to=10)
-        value_mha = op.Cast(value, to=10)
+        # PackedMHA supports float32 and float16 only; cast bfloat16 builds to
+        # float16 and leave float32/float16 native to preserve precision.
+        if get_build_dtype() == ir.DataType.BFLOAT16:
+            query_mha = op.Cast(query, to=ir.DataType.FLOAT16)
+            key_mha = op.Cast(key, to=ir.DataType.FLOAT16)
+            value_mha = op.Cast(value, to=ir.DataType.FLOAT16)
+        else:
+            query_mha, key_mha, value_mha = query, key, value
 
         # Emit PackedMultiHeadAttention
         attn_out = op.PackedMultiHeadAttention(

src/mobius/components/_qwen25_vl_vision_test.py

@@ -0,0 +1,83 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Graph-construction tests for the Qwen2.5-VL vision encoder packed path.
+
+These exercise the EP-gated ``_emit_packed_mha`` branch (only taken when the
+active EP advertises ``supports_packed_multi_head_attention``, e.g. CUDA),
+which the default CPU build does not reach.
+"""
+
+from __future__ import annotations
+
+import onnx_ir as ir
+
+from mobius._build_context import build_context
+from mobius._execution_providers import ep_registry
+from mobius._testing import count_op_type, create_test_builder, create_test_input
+from mobius.components._qwen25_vl_vision import Qwen25VLVisionModel
+
+_PATCH_DIM = 3 * 2 * 14 * 14  # in_channels * temporal_patch * patch * patch
+
+
+def _build_vision_graph(dtype: ir.DataType) -> ir.Graph:
+    """Build the Qwen2.5-VL vision encoder graph under the CUDA EP.
+
+    ``fullatt_block_indexes=[1]`` makes block 0 windowed and block 1 full, so
+    both attention variants flow through the packed path.
+    """
+    module = Qwen25VLVisionModel(
+        depth=2,
+        hidden_size=64,
+        intermediate_size=128,
+        num_heads=2,
+        patch_size=14,
+        temporal_patch_size=2,
+        in_channels=3,
+        out_hidden_size=64,
+        spatial_merge_size=2,
+        fullatt_block_indexes=[1],
+        window_size=112,
+    )
+    builder, op, graph = create_test_builder()
+    pixel_values = create_test_input(builder, "pixel_values", ["N", _PATCH_DIM], dtype=dtype)
+    grid = create_test_input(
+        builder, "image_grid_thw", ["num_images", 3], dtype=ir.DataType.INT64
+    )
+    with build_context(ep_registry.require("cuda"), dtype):
+        out = module(op, pixel_values, grid)
+    out.name = "image_features"
+    graph.outputs.append(out)
+    return graph
+
+
+def _count_cast_to(graph: ir.Graph, target: ir.DataType) -> int:
+    count = 0
+    for node in graph:
+        if node.op_type == "Cast" and int(node.attributes["to"].value) == int(target):
+            count += 1
+    return count
+
+
+class TestQwen25VLVisionPackedPath:
+    def test_cuda_emits_packed_mha_with_helper_token_offset(self):
+        graph = _build_vision_graph(ir.DataType.FLOAT)
+        # One PackedMultiHeadAttention per transformer block, no standard
+        # Attention fallback on CUDA.
+        assert count_op_type(graph, "PackedMultiHeadAttention") == 2
+        assert count_op_type(graph, "Attention") == 0
+        # token_offset comes from build_packed_token_offset (Compress-based
+        # valid/padding split), not the old (1, N) identity Range.
+        assert count_op_type(graph, "Compress") >= 4  # 2 blocks x (valid + padding)
+
+    def test_float32_build_keeps_native_dtype(self):
+        # f32 is supported natively by the kernel: no down-cast to float16.
+        graph = _build_vision_graph(ir.DataType.FLOAT)
+        assert _count_cast_to(graph, ir.DataType.FLOAT16) == 0
+
+    def test_bfloat16_build_casts_qkv_to_float16(self):
+        # bf16 is unsupported by the kernel, so q/k/v are cast to float16:
+        # 3 casts (q, k, v) per block x 2 blocks.
+        graph = _build_vision_graph(ir.DataType.BFLOAT16)
+        assert count_op_type(graph, "PackedMultiHeadAttention") == 2
+        assert _count_cast_to(graph, ir.DataType.FLOAT16) == 6

src/mobius/components/_qwen3_vl_vision.py

@@ -25,8 +25,8 @@
 import onnx_ir as ir
 from onnxscript import OpBuilder, nn
 
-from mobius._build_context import ep_capabilities
-from mobius.components._common import LayerNorm, Linear
+from mobius._build_context import ep_capabilities, get_build_dtype
+from mobius.components._common import LayerNorm, Linear, build_packed_token_offset
 from mobius.components._mlp import FCMLP
 from mobius.components._scan_utils import (
     compact_scan_output,
@@ -219,34 +219,35 @@ def _emit_packed_mha(self, op, query, key, value, cu_seqlens, hidden_states):
         """Emit com.microsoft.PackedMultiHeadAttention.
 
         Uses cu_seqlens natively, avoiding the O(N^2) block-diagonal bias.
-        PackedMHA supports float32/float16 only; bf16 inputs are cast to f16.
+        Each sub-sequence delimited by ``cu_seqlens`` (one per image/frame)
+        is one packed batch element, so block-diagonal masking is expressed
+        through the varlen contract rather than an explicit bias.
 
         Args:
             query, key: (total_seq, hidden_size) after rotary embedding
-            value: (total_seq, 3 * hidden_size) — full QKV output, need V only
-            cu_seqlens: (num_sub_seqs + 1,) INT32/INT64
+            value: (total_seq, hidden_size) from QKV split
+            cu_seqlens: (num_sub_seqs + 1,) cumulative sequence lengths
             hidden_states: original input, used only for shape
         """
-        total_seq = op.Shape(hidden_states, start=0, end=1)
-        total_seq_scalar = op.Squeeze(total_seq)
-
-        # token_offset: identity mapping for packed (no-padding) input.
-        token_offset = op.Unsqueeze(
-            op.Range(
-                op.Constant(value_int=0),
-                total_seq_scalar,
-                op.Constant(value_int=1),
-            ),
-            [0],
-        )
-        token_offset = op.Cast(token_offset, to=6)  # INT32
+        # token_offset: (num_sub_seqs, max_seq_len) mapping packed tokens to
+        # their padded (batch, seq) layout. ORT derives batch_size from
+        # token_offset.shape[0] and requires cumulative_sequence_length to have
+        # length batch_size + 1, so this MUST encode every sub-sequence (image
+        # or frame), not a single (1, N) batch — otherwise the kernel rejects
+        # multi-sequence cu_seqlens and computes full instead of block-diagonal
+        # attention.
+        token_offset = build_packed_token_offset(op, cu_seqlens)
 
         cu_seqlens_int32 = op.Cast(cu_seqlens, to=6)  # INT32
 
-        # PackedMHA doesn't support bfloat16; cast to float16 if needed
-        query_mha = op.Cast(query, to=10)  # FLOAT16
-        key_mha = op.Cast(key, to=10)
-        value_mha = op.Cast(value, to=10)
+        # PackedMHA supports float32 and float16 only; cast bfloat16 builds to
+        # float16 and leave float32/float16 native to preserve precision.
+        if get_build_dtype() == ir.DataType.BFLOAT16:
+            query_mha = op.Cast(query, to=ir.DataType.FLOAT16)
+            key_mha = op.Cast(key, to=ir.DataType.FLOAT16)
+            value_mha = op.Cast(value, to=ir.DataType.FLOAT16)
+        else:
+            query_mha, key_mha, value_mha = query, key, value
 
         attn_output = op.PackedMultiHeadAttention(
             query_mha,