[NPU] Enable Qwen3-VL GRPO (single- & multi-turn) on Ascend

docker/npu_patch/megatron-bridge.patch

@@ -5,6 +5,7 @@ index a0421273..7f9203ab 100644
 @@ -1097,15 +1097,19 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
              "LinearCrossEntropyModule",
              "TEColumnParallelLinear",
++            "MindSpeedTEColumnParallelLinear",
              "TELayerNormColumnParallelLinear",
 +            "MindSpeedTELayerNormColumnParallelLinear",
              "TEColumnParallelGroupedLinear",
@@ -298,3 +299,77 @@ index 60cc091c..1dd5c808 100644
  from megatron.core.models.gpt.gpt_layer_specs import (
      get_gpt_decoder_block_spec,
      get_gpt_layer_local_spec,
+diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+index 7775c11c..c7b094cf 100644
+--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
++++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/model.py
+@@ -26,6 +26,7 @@ from megatron.core.packed_seq_params import PackedSeqParams
+ from megatron.core.process_groups_config import ProcessGroupCollection
+ from megatron.core.transformer import MegatronModule
+ from megatron.core.transformer.spec_utils import ModuleSpec
++from megatron.core.utils import nvtx_range_pop, nvtx_range_push
+ from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig as Qwen3VLConfigHF
+
+ from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.attention import Qwen3VLSelfAttention
+@@ -294,7 +295,7 @@ class Qwen3VLModel(MegatronModule):
+         # position ids is computed within the model
+         position_ids = None
+
+-        torch.cuda.nvtx.range_push("Qwen3VLModel.forward.pre_process")
++        nvtx_range_push(msg="Qwen3VLModel.forward.pre_process")
+
+         cp_rank = self.pg_collection.cp.rank()
+         cp_size = self.pg_collection.cp.size()
+@@ -387,7 +388,7 @@ class Qwen3VLModel(MegatronModule):
+                 combined_embeddings = split_data_cp_rank(combined_embeddings, cp_size, 0, cp_rank)
+             if packed_seq_params is not None:
+                 if attention_mask is None:
+-                    attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)
++                    attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device)
+                 input_ids_thd, _ = preprocess_packed_seqs(
+                     input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
+                 )
+@@ -443,7 +444,7 @@ class Qwen3VLModel(MegatronModule):
+             # convert lm_input_ids to THD format so it matches position_ids.
+             if packed_seq_params is not None:
+                 if attention_mask is None:
+-                    attention_mask = torch.ones_like(input_ids, dtype=torch.int32, device=input_ids.device)
++                    attention_mask = torch.ones_like(input_ids, dtype=torch.bool, device=input_ids.device)
+                 lm_input_ids, _ = preprocess_packed_seqs(
+                     input_ids, attention_mask, pre_process=True, pg_collection=self.pg_collection
+                 )
+@@ -499,8 +500,8 @@ class Qwen3VLModel(MegatronModule):
+                 attention_mask = None
+                 self.language_model.rotary_pos_emb.is_thd_format = True
+
+-        torch.cuda.nvtx.range_pop()
+-        torch.cuda.nvtx.range_push("Qwen3VLModel.forward.language_model")
++        nvtx_range_pop(msg="Qwen3VLModel.forward.pre_process")
++        nvtx_range_push(msg="Qwen3VLModel.forward.language_model")
+
+         output = self.language_model(
+             input_ids=lm_input_ids,
+@@ -516,6 +517,6 @@ class Qwen3VLModel(MegatronModule):
+             **(extra_block_kwargs or {}),
+             **kwargs,
+         )
+-        torch.cuda.nvtx.range_pop()
++        nvtx_range_pop(msg="Qwen3VLModel.forward.language_model")
+
+         return output
+diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
+index b714d0f7..42fcab74 100644
+--- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
++++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/utils.py
+@@ -668,6 +668,11 @@ def preprocess_packed_seqs(
+     """
+     batch_size = input_ids.shape[0]
+
++    # Ensure boolean dtype for correct advanced indexing (bool → mask select,
++    # int → fancy index which silently corrupts data when values are 0/1).
++    if attention_mask.dtype != torch.bool:
++        attention_mask = attention_mask.bool()
++
+     seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+     if pg_collection is not None:
+         tp_size = pg_collection.tp.size()

docker/npu_patch/mindspeed.patch

@@ -98,7 +98,7 @@ index f14b231d..4615590e 100644
          setattr(full_args, k, v)
      MindSpeedFeaturesManager.apply_features_pre_patches(full_args)
 diff --git a/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py b/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
-index ac4eabe5..78c5866d 100644
+index ac4eabe5..061377cc 100644
 --- a/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
 +++ b/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
 @@ -330,6 +330,8 @@ class DotProductAttention(torch.nn.Module):
@@ -110,18 +110,52 @@ index ac4eabe5..78c5866d 100644
      ) -> torch.Tensor:
          """
          Dot Product Attention Layer.
-@@ -659,7 +661,9 @@ class MindSpeedTEDotProductAttention(DotProductAttention):
+@@ -653,18 +655,39 @@ class MindSpeedTEDotProductAttention(DotProductAttention):
+         packed_seq_params: PackedSeqParams = None,
+         ):
+         """Forward."""
++        packed_seq_kwargs = (
++            {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
++            if packed_seq_params is not None
++            else {}
++        )
++
++        # Honor the per-call attn_mask_type. Megatron passes AttnMaskType.no_mask for the
++        # (bidirectional) vision tower and AttnMaskType.causal for the causal LM decoder.
++        # The previous code unconditionally built a triu(2048) causal mask and forwarded
++        # config.attention_mask_type, forcing CAUSAL attention even for no_mask -> the
++        # Qwen3-VL vision tower's full attention was computed causally (wrong image embeds).
++        if attn_mask_type == AttnMaskType.no_mask:
++            # Full (bidirectional) attention: no causal mask; per-image segmentation is
++            # handled by cu_seqlens in packed_seq_params (sparse_mode=0 via 'no_mask').
++            core_attn_out = super().forward(
++                query,
++                key,
++                value,
++                None,
++                attn_mask_type='no_mask',
++                **packed_seq_kwargs,
++            )
++            return core_attn_out
++
+         if (
+             attention_mask is None and
+             self.attn_mask_type == AttnMaskType.causal
          ) and not getattr(self.config, 'is_llava', False):
              self.config.sparse_mode = 2
              attention_mask = get_attention_mask(self.config)
 -
+-        packed_seq_kwargs = (
+-            {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
+-            if packed_seq_params is not None
+-            else {}
+-        )
 +        attention_mask = torch.triu(
 +            torch.ones((2048, 2048),
-+                       device=query.device, dtype=torch.bool), diagonal=1) 
-         packed_seq_kwargs = (
-             {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
-             if packed_seq_params is not None
-
++                       device=query.device, dtype=torch.bool), diagonal=1)
+
+         core_attn_out = super().forward(
+             query,
 diff --git a/mindspeed/core/fusions/fused_rope.py b/mindspeed/core/fusions/fused_rope.py
 index 70f7cb08..15189c2d 100644
 --- a/mindspeed/core/fusions/fused_rope.py