fix num_samples attributes error

src/mcore_bridge/model/gpts/qwen3_next.py

@@ -20,7 +20,7 @@
 
 from mcore_bridge.bridge import GPTBridge
 from mcore_bridge.config import ModelConfig
-from mcore_bridge.utils import get_env_args, get_local_layer_specs, get_logger
+from mcore_bridge.utils import get_env_args, get_local_layer_specs, get_logger, get_num_samples
 
 from ..constant import ModelType
 from ..register import ModelLoader, ModelMeta, register_model
@@ -470,12 +470,12 @@ def forward(self, hidden_states: torch.Tensor, **kwargs):
         # Note: for packed inputs, we do not perform padding_free unpadding.
         # Doing so would allow different sequences to see each other; for efficiency we keep this implementation.
         if thd_format:
+            num_samples = get_num_samples(packed_seq_params)
             max_seqlen_q = int(packed_seq_params.max_seqlen_q)
-            new_hidden_states = hidden_states.new_zeros(
-                (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1]))
-            attention_mask = hidden_states.new_zeros((packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool)
+            new_hidden_states = hidden_states.new_zeros((num_samples, max_seqlen_q, hidden_states.shape[-1]))
+            attention_mask = hidden_states.new_zeros((num_samples, max_seqlen_q), dtype=torch.bool)
             cu_seqlens_q = packed_seq_params.cu_seqlens_q
-            for i in range(packed_seq_params.num_samples):
+            for i in range(num_samples):
                 start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1]
                 attention_mask[i, :end - start] = True
                 new_hidden_states[i, :end - start] = hidden_states[start:end, 0]

src/mcore_bridge/model/mm_gpts/qwen3_5.py

@@ -7,7 +7,7 @@
 from megatron.core.transformer.attention import SelfAttentionSubmodules
 from megatron.core.transformer.transformer_config import TransformerConfig
 
-from mcore_bridge.utils import get_env_args
+from mcore_bridge.utils import get_env_args, get_num_samples
 
 from ..constant import ModelType
 from ..gpts.qwen3_next import Qwen3NextBridge, Qwen3NextLoader, resolve_gdn_attention_mask
@@ -40,12 +40,12 @@ def forward(self, hidden_states: torch.Tensor, **kwargs):
         # Note: for packed inputs, we do not perform padding_free unpadding.
         # Doing so would allow different sequences to see each other; for efficiency we keep this implementation.
         if thd_format:
+            num_samples = get_num_samples(packed_seq_params)
             max_seqlen_q = int(packed_seq_params.max_seqlen_q)
-            new_hidden_states = hidden_states.new_zeros(
-                (packed_seq_params.num_samples, max_seqlen_q, hidden_states.shape[-1]))
-            attention_mask = hidden_states.new_zeros((packed_seq_params.num_samples, max_seqlen_q), dtype=torch.bool)
+            new_hidden_states = hidden_states.new_zeros((num_samples, max_seqlen_q, hidden_states.shape[-1]))
+            attention_mask = hidden_states.new_zeros((num_samples, max_seqlen_q), dtype=torch.bool)
             cu_seqlens_q = packed_seq_params.cu_seqlens_q
-            for i in range(packed_seq_params.num_samples):
+            for i in range(num_samples):
                 start, end = cu_seqlens_q[i], cu_seqlens_q[i + 1]
                 attention_mask[i, :end - start] = True
                 new_hidden_states[i, :end - start] = hidden_states[start:end, 0]

src/mcore_bridge/model/modules/ple.py

@@ -9,7 +9,7 @@
 from torch import nn
 from typing import List, Optional
 
-from ...utils.megatron_utils import reconstruct_tensor_cp, split_cp_inputs
+from ...utils.megatron_utils import get_num_samples, reconstruct_tensor_cp, split_cp_inputs
 from .hyper_connection_gated import Qwen4ExpTextGroupedRMSNorm
 
 _MASK64 = (1 << 64) - 1
@@ -382,7 +382,7 @@ def _forward_impl(
         """hidden_states: [s, b, nH] (bsh) or thd [T, 1, nH]; input_ids: [b, s] or [1, T]."""
         thd = packed_seq_params is not None and getattr(packed_seq_params, 'qkv_format', 'bshd') == 'thd'
         if thd:
-            num_samples = packed_seq_params.num_samples
+            num_samples = get_num_samples(packed_seq_params)
             # PackedSeqParams.max_seqlen_q is declared `int` in mcore and swift
             # normalizes it to int, so `.item()` would raise AttributeError;
             # tolerate a 0-d tensor from other callers.

src/mcore_bridge/utils/__init__.py

@@ -3,8 +3,8 @@
 from .env import get_dist_setting, get_node_setting, is_dist, is_last_rank, is_local_master, is_master
 from .import_utils import _LazyModule, is_flash_attn_3_available
 from .logger import get_logger
-from .megatron_utils import (get_local_layer_specs, reconstruct_tensor_cp, roll_tensor, set_random_seed,
-                             split_cp_inputs, unwrap_model)
+from .megatron_utils import (get_local_layer_specs, get_num_samples, reconstruct_tensor_cp, roll_tensor,
+                             set_random_seed, split_cp_inputs, unwrap_model)
 from .safetensors import SafetensorLazyLoader, StreamingSafetensorSaver
 from .torch_utils import gc_collect, get_current_device, safe_ddp_context, to_device
 from .utils import deep_getattr, get_env_args, json_parse_to_dict, patch_deepcopy

src/mcore_bridge/utils/megatron_utils.py

@@ -83,6 +83,30 @@ def split_cp_inputs(inputs: torch.Tensor,
     return torch.cat(new_inputs, dim=dim)
 
 
+def get_num_samples(packed_seq_params) -> int:
+    """Number of real (unpadded) samples packed into a THD micro-batch.
+
+    ``cu_seqlens_q.numel() - 1`` is *not* this number: the data pipeline pads the packed
+    sequence (see swift ``get_padding_to``) and the padding tokens carry ``position_ids == 0``,
+    so the padded tail opens extra segments (one per token when cp_size == 1, one per
+    ``2 * cp_size`` tokens otherwise). Using the segment count would inflate the
+    ``[num_samples, max_seqlen_q, h]`` buffers built by the GDN/PLE THD paths by up to
+    ``padding_to - 1`` extra rows.
+
+    ``num_samples``/``seq_lens`` are attached to ``PackedSeqParams`` by the data pipeline;
+    they are not mcore dataclass fields. Fall back to the segment count only when neither is
+    available, i.e. for callers that build ``PackedSeqParams`` directly (the convert-precision
+    check and the README example), which run tiny batches.
+    """
+    num_samples = getattr(packed_seq_params, 'num_samples', None)
+    if num_samples is not None:
+        return int(num_samples)
+    seq_lens = getattr(packed_seq_params, 'seq_lens', None)
+    if seq_lens is not None:
+        return int(seq_lens.shape[0])
+    return int(packed_seq_params.cu_seqlens_q.numel()) - 1
+
+
 def reconstruct_tensor_cp(tensor, packed_seq_params, dim: int) -> torch.Tensor:
     """In CP mode, all-gather and undo the load-balanced (zigzag) chunking
     produced by ``split_cp_inputs``, restoring the full sequence in original