diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/attention.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/attention.py index fb268a0bae..4965ecf108 100644 --- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/attention.py +++ b/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: diff --git a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_attention.py b/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_attention.py index ba6d4ded21..5e63cc4ca2 100644 --- a/tests/unit_tests/models/qwen_vl/modelling_qwen3_vl/test_attention.py +++ b/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: