fix(models): use physical boundaries for Qwen packed attention

src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/attention.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from copy import copy
+
 import torch
 from einops import rearrange
 from megatron.core.extensions.transformer_engine import TEDotProductAttention
@@ -56,6 +58,39 @@ def _qwen_attention_mask_for_core_attention(
     return (~valid_mask).unsqueeze(1).unsqueeze(1)
 
 
+def _get_core_attention_packed_seq_params(
+    packed_seq_params: PackedSeqParams | None,
+) -> PackedSeqParams | None:
+    """Use physical THD boundaries for Qwen core attention.
+
+    Qwen's hybrid layers need both logical and physical boundaries outside
+    attention, but TE attention kernels consume the full physical token stream.
+    Passing only the physical boundaries avoids relying on TE's padded-layout
+    inference, which ignores a difference at the final cumulative offset.
+
+    Args:
+        packed_seq_params: Packed metadata shared by the Qwen decoder.
+
+    Returns:
+        A shallow copy using physical q/kv boundaries for core attention, or
+        the original value when no padded layout is present.
+    """
+    if packed_seq_params is None or packed_seq_params.cu_seqlens_q_padded is None:
+        return packed_seq_params
+
+    core_attention_params = copy(packed_seq_params)
+    core_attention_params.cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded
+    core_attention_params.cu_seqlens_kv = (
+        packed_seq_params.cu_seqlens_kv_padded
+        if packed_seq_params.cu_seqlens_kv_padded is not None
+        else packed_seq_params.cu_seqlens_q_padded
+    )
+    core_attention_params.cu_seqlens_q_padded = None
+    core_attention_params.cu_seqlens_kv_padded = None
+    core_attention_params.pad_between_seqs = False
+    return core_attention_params
+
+
 class Qwen3VLSelfAttention(SelfAttention):
     """
     Overrides the SelfAttention class, the difference is that qwen3vl uses apply_rotary_pos_emb_absolute
@@ -264,6 +299,7 @@ def forward(
         # ==================================
 
         nvtx_range_push(suffix="core_attention")
+        core_attention_packed_seq_params = _get_core_attention_packed_seq_params(packed_seq_params)
         core_attention_mask = _qwen_attention_mask_for_core_attention(
             self.core_attention,
             attention_mask,
@@ -285,7 +321,7 @@ def forward(
                 core_attention_mask,
                 attn_mask_type=attn_mask_type,
                 attention_bias=attention_bias,
-                packed_seq_params=packed_seq_params,
+                packed_seq_params=core_attention_packed_seq_params,
             )
         else:
             if inference_context is None or inference_context.is_static_batching():
@@ -297,7 +333,7 @@ def forward(
                     core_attention_mask,
                     attn_mask_type=attn_mask_type,
                     attention_bias=attention_bias,
-                    packed_seq_params=packed_seq_params,
+                    packed_seq_params=core_attention_packed_seq_params,
                 )
 
             else:

tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_attention.py

@@ -25,12 +25,53 @@
 import torch.distributed as dist
 from megatron.core import parallel_state
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
+from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 from megatron.core.transformer import TransformerConfig
 from megatron.core.transformer.attention import AttnMaskType
 
-from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention
+from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import (
+    Qwen3VLSelfAttention,
+    _get_core_attention_packed_seq_params,
+)
+
+
+def test_core_attention_packed_params_use_physical_tail_boundary():
+    logical = torch.tensor([0, 5], dtype=torch.int32)
+    physical = torch.tensor([0, 8], dtype=torch.int32)
+    packed_seq_params = PackedSeqParams(
+        qkv_format="thd",
+        cu_seqlens_q=logical,
+        cu_seqlens_kv=logical,
+        cu_seqlens_q_padded=physical,
+        cu_seqlens_kv_padded=physical,
+        max_seqlen_q=8,
+        max_seqlen_kv=8,
+    )
+
+    actual = _get_core_attention_packed_seq_params(packed_seq_params)
+
+    assert actual is not packed_seq_params
+    assert actual.cu_seqlens_q is physical
+    assert actual.cu_seqlens_kv is physical
+    assert actual.cu_seqlens_q_padded is None
+    assert actual.cu_seqlens_kv_padded is None
+    assert actual.pad_between_seqs is False
+    assert packed_seq_params.cu_seqlens_q is logical
+    assert packed_seq_params.cu_seqlens_q_padded is physical
+
+
+def test_core_attention_packed_params_preserve_unpadded_metadata():
+    packed_seq_params = PackedSeqParams(
+        qkv_format="thd",
+        cu_seqlens_q=torch.tensor([0, 8], dtype=torch.int32),
+        cu_seqlens_kv=torch.tensor([0, 8], dtype=torch.int32),
+        max_seqlen_q=8,
+        max_seqlen_kv=8,
+    )
+
+    assert _get_core_attention_packed_seq_params(packed_seq_params) is packed_seq_params
 
 
 class TestQwen3VLSelfAttention: