[https://nvbugs/6562942][fix] Load Qwen3.5 packed INT4 checkpoints

tensorrt_llm/_torch/models/checkpoints/hf/qwen3_5_weight_mapper.py

@@ -12,14 +12,30 @@
 from tensorrt_llm._torch.models.checkpoints.hf.qwen3_next_weight_mapper import (
     Qwen3NextHfWeightMapper,
 )
-from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm._torch.models.modeling_utils import (
+    concatenate_weights_by_tp_rank,
+    register_mapper,
+)
 from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode
 from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner
 from tensorrt_llm.quantization import QuantAlgo
 
 _FP8_2D_BLOCK_SIZE = 128
 
 
+def _check_packed_int4(
+    tensors: dict[str, torch.Tensor], keys: set[str], rows: tuple[int, ...], prefix: str
+) -> tuple[int, ...]:
+    """Validate packed projection storage and return physical row counts."""
+    if not all(tensors[key].dtype in (torch.uint8, torch.int8) for key in keys):
+        raise ValueError(f"Cannot fuse packed INT4 and unpacked weights for {prefix}")
+    if len({tensors[key].dtype for key in keys}) != 1:
+        raise ValueError(f"Packed INT4 projections must have the same storage dtype for {prefix}")
+    if any(row % 2 != 0 for row in rows):
+        raise ValueError(f"INT4 projection dimensions must be even for {prefix}")
+    return tuple(row // 2 for row in rows)
+
+
 @register_mapper("HF", "Qwen3_5ForConditionalGeneration")
 @register_mapper("HF", "QwenImageBenchForConditionalGeneration")
 @register_mapper("HF", "Qwen3_5MoeForCausalLM")
@@ -43,7 +59,9 @@ class Qwen3_5MoeHfWeightMapper(Qwen3NextHfWeightMapper):
        Qwen3Next checkpoints store pre-packed in_proj_qkvz and in_proj_ba
        tensors.  Qwen3.5 checkpoints store them as separate in_proj_qkv + z
        (or fully split q/k/v/z) and b + a tensors.  This mapper packs them
-       into the grouped-interleaved layout that TRT-LLM expects.
+       into the grouped-interleaved layout consumed by the Qwen3Next mapper.
+       INT4 projections are concatenated directly in per-rank dense order,
+       preserving the two output channels stored in each byte.
        For FP8_BLOCK_SCALES checkpoints, the packed qkvz tensor is then
        dequantized to bf16 as a temporary workaround for TP loading
        (handled in _dequantize_linear_attn_fp8_qkvz).  For MIXED_PRECISION
@@ -483,7 +501,7 @@ def _dequantize_linear_attn_fp8_per_tensor(self, weights: dict) -> dict:
             if name not in drop_keys
         }
 
-    def _pack_split_projections(self, weights: dict) -> dict:
+    def _pack_split_projections(self, weights: dict, dense_layout: bool = False) -> dict:
         config = self.config.pretrained_config
         num_k_groups = config.linear_num_key_heads
         num_v_heads = config.linear_num_value_heads
@@ -492,6 +510,13 @@ def _pack_split_projections(self, weights: dict) -> dict:
             f"linear_num_key_heads ({num_k_groups})"
         )
 
+        pack_projection_tensor = self._pack_projection_tensor
+        num_groups = num_k_groups
+        if dense_layout:
+            mapping = self.config.mapping
+            num_groups = 1 if mapping.enable_attention_dp else mapping.tp_size
+            pack_projection_tensor = concatenate_weights_by_tp_rank
+
         grouped_weights = defaultdict(dict)
         packed_weights = {}
 
@@ -508,9 +533,43 @@ def _pack_split_projections(self, weights: dict) -> dict:
         expected_ba = config.linear_num_value_heads
 
         for (prefix, suffix), tensors in grouped_weights.items():
+            if dense_layout and suffix == "pre_quant_scale":
+                # Input scales are shared by fused projections, not concatenated by output row.
+                for projection, keys in (
+                    ("qkvz", ("qkv", "z") if "qkv" in tensors else ("q", "k", "v", "z")),
+                    ("ba", ("b", "a")),
+                ):
+                    if not set(keys) & tensors.keys():
+                        continue
+                    if missing := set(keys) - tensors.keys():
+                        raise ValueError(
+                            f"Cannot fuse {prefix}.in_proj_{projection}.pre_quant_scale: "
+                            f"missing scales for {sorted(missing)}"
+                        )
+                    scales = [
+                        tensors[key]
+                        if isinstance(tensors[key], torch.Tensor)
+                        else tensors[key][...]
+                        for key in keys
+                    ]
+                    if any(scale.shape != (config.hidden_size,) for scale in scales):
+                        raise ValueError(
+                            f"Expected pre_quant_scale shape ({config.hidden_size},) for {prefix}"
+                        )
+                    if not all(torch.equal(scales[0], scale) for scale in scales[1:]):
+                        raise ValueError(
+                            f"Cannot fuse {prefix}.in_proj_{projection}: pre_quant_scale vectors "
+                            "must be identical; export with shared input scales"
+                        )
+                    packed_name = f"{prefix}.in_proj_{projection}.{suffix}"
+                    if packed_name in packed_weights:
+                        raise ValueError(f"Packed projection {packed_name} already exists")
+                    packed_weights[packed_name] = scales[0]
+                continue
             # `weight_scale_inv` (loaded by FP8BlockScalesLinearMethod) is the
             # only suffix stored in 2D-block format [ceil(out/block_size), ceil(in/block_size)];
-            # weight/bias/weight_scale all keep out_features as their leading dim.
+            # INT4 weights pack two output channels per byte; their scales and
+            # biases retain the logical output dimension.
             if suffix == "weight_scale_inv":
                 row_q = math.ceil(expected_q / _FP8_2D_BLOCK_SIZE)
                 row_v = math.ceil(expected_v / _FP8_2D_BLOCK_SIZE)
@@ -522,13 +581,22 @@ def _pack_split_projections(self, weights: dict) -> dict:
 
             qkvz_keys = {"qkv", "q", "k", "v", "z"} & tensors.keys()
             if qkvz_keys:
+                if (
+                    dense_layout
+                    and suffix == "weight"
+                    and any(tensors[key].dtype in (torch.uint8, torch.int8) for key in qkvz_keys)
+                ):
+                    row_q, row_v = _check_packed_int4(tensors, qkvz_keys, (row_q, row_v), prefix)
                 if "qkv" in tensors:
                     missing = {"qkv", "z"} - tensors.keys()
                     assert not missing, (
                         f"Missing split projections {sorted(missing)} for {prefix}.{suffix}"
                     )
+                    split_q, split_v = (
+                        (expected_q, expected_v) if suffix == "weight_scale_inv" else (row_q, row_v)
+                    )
                     q_tensor, k_tensor, v_tensor = split_packed_qkv(
-                        tensors["qkv"], expected_q, expected_v
+                        tensors["qkv"], split_q, split_v
                     )
                 else:
                     missing = {"q", "k", "v", "z"} - tensors.keys()
@@ -548,8 +616,8 @@ def _pack_split_projections(self, weights: dict) -> dict:
                 assert packed_name not in packed_weights, (
                     f"Packed projection {packed_name} already exists"
                 )
-                packed_weights[packed_name] = self._pack_projection_tensor(
-                    [q_tensor, k_tensor, v_tensor, tensors["z"]], num_k_groups
+                packed_weights[packed_name] = pack_projection_tensor(
+                    [q_tensor, k_tensor, v_tensor, tensors["z"]], num_groups
                 )
 
             ba_keys = {"b", "a"} & tensors.keys()
@@ -560,15 +628,19 @@ def _pack_split_projections(self, weights: dict) -> dict:
                 )
 
                 if suffix == "weight":
+                    if dense_layout and any(
+                        tensors[key].dtype in (torch.uint8, torch.int8) for key in ba_keys
+                    ):
+                        (row_ba,) = _check_packed_int4(tensors, ba_keys, (row_ba,), prefix)
                     assert tensors["b"].shape[0] == row_ba
                     assert tensors["a"].shape[0] == row_ba
 
                 packed_name = f"{prefix}.in_proj_ba.{suffix}"
                 assert packed_name not in packed_weights, (
                     f"Packed projection {packed_name} already exists"
                 )
-                packed_weights[packed_name] = self._pack_projection_tensor(
-                    [tensors["b"], tensors["a"]], num_k_groups
+                packed_weights[packed_name] = pack_projection_tensor(
+                    [tensors["b"], tensors["a"]], num_groups
                 )
 
         return packed_weights
@@ -694,16 +766,26 @@ def preprocess_weights(self, weights: dict, allow_partial_loading: bool = False)
         if quant_algo in (QuantAlgo.MIXED_PRECISION, QuantAlgo.NVFP4):
             normalized_weights = self._dequantize_lm_head_nvfp4(normalized_weights)
 
-        packed_weights = self._pack_split_projections(normalized_weights)
+        dense_layout = quant_algo == QuantAlgo.W4A16_AWQ
+        if dense_layout:
+            # Apply MTP renaming and non-projection transforms while in_proj
+            # tensors are still split. The inherited mapper passes those
+            # through, so the directly fused result is not permuted again.
+            normalized_weights = super().preprocess_weights(
+                normalized_weights, allow_partial_loading=allow_partial_loading
+            )
+        packed_weights = self._pack_split_projections(normalized_weights, dense_layout=dense_layout)
         if quant_algo == QuantAlgo.FP8_BLOCK_SCALES and not is_modelopt_pb_wo:
             packed_weights = self._dequantize_linear_attn_fp8_qkvz(packed_weights)
 
         if not getattr(self.config.pretrained_config, "num_experts", 0):
             packed_weights = self._remap_dense_mlp_weights(packed_weights)
 
-        processed_weights = super().preprocess_weights(
-            packed_weights, allow_partial_loading=allow_partial_loading
-        )
+        processed_weights = packed_weights
+        if not dense_layout:
+            processed_weights = super().preprocess_weights(
+                packed_weights, allow_partial_loading=allow_partial_loading
+            )
         if is_consumable and not isinstance(processed_weights, ConsumableWeightsDict):
             return ConsumableWeightsDict(processed_weights)
         return processed_weights

tensorrt_llm/_torch/models/checkpoints/hf/qwen4_exp_weight_mapper.py

@@ -76,7 +76,10 @@
 from torch import nn
 
 from tensorrt_llm._torch.models.checkpoints.hf.qwen2_moe_weight_mapper import Qwen2MoeHfWeightMapper
-from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm._torch.models.modeling_utils import (
+    concatenate_weights_by_tp_rank,
+    register_mapper,
+)
 from tensorrt_llm._torch.modules.qwen4_exp.ple import Qwen4ExpPinnedHostEmbedding
 from tensorrt_llm._torch.moe.fused_moe.interface import MoEWeightLoadingMode
 from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner
@@ -89,36 +92,6 @@
 _PER_EXPERT_PROJECTION_PATTERN = re.compile(r"^\d+\.(?:gate_proj|up_proj|down_proj)\.")
 
 
-def _rank_block(components: list[torch.Tensor], tp_size: int) -> torch.Tensor:
-    """Concat `components` (each `[out_i, ...]`) into the per-rank column-
-    parallel row order a contiguous `TensorParallelMode.COLUMN` split recovers.
-
-    At `tp_size == 1` this is a plain `[c0 | c1 | ...]` concat. At
-    `tp_size > 1` each component is split across ranks and the rank-`r` slices
-    are grouped together (`[c0_r0 c1_r0 ... | c0_r1 c1_r1 ... | ...]`), so a
-    later contiguous per-rank row split hands each rank its own head slice of
-    every component.
-    """
-    if not components:
-        raise ValueError("rank blocking requires at least one projection")
-    if tp_size <= 0:
-        raise ValueError(f"tp_size must be positive, got {tp_size}")
-    trailing_shape = components[0].shape[1:]
-    for component in components:
-        if component.ndim == 0 or component.shape[1:] != trailing_shape:
-            raise ValueError("rank-blocked projections must have matching trailing dimensions")
-        if component.shape[0] % tp_size:
-            raise ValueError(
-                f"projection rows {component.shape[0]} are not divisible by tp_size={tp_size}"
-            )
-    if tp_size == 1:
-        return torch.cat(components, dim=0)
-    rows: list[torch.Tensor] = []
-    for rank in range(tp_size):
-        rows.extend(split(c, tp_size, rank) for c in components)
-    return torch.cat(rows, dim=0)
-
-
 def _normalize_moe_module_weights(
     module_weights: dict[str, torch.Tensor], config
 ) -> tuple[dict[str, torch.Tensor], MoEWeightLoadingMode]:
@@ -363,7 +336,7 @@ def preprocess_weights(self, weights: dict, allow_partial_loading: bool = False)
                         f"{key} convolution width is {w.shape[1]}, expected {conv_kernel}"
                     )
                 conv_q, conv_k, conv_v = torch.split(w, [key_dim, key_dim, value_dim], dim=0)
-                new_weights[key] = _rank_block([conv_q, conv_k, conv_v], tp_size)
+                new_weights[key] = concatenate_weights_by_tp_rank([conv_q, conv_k, conv_v], tp_size)
             else:
                 new_weights[key] = tensor
 
@@ -376,7 +349,7 @@ def preprocess_weights(self, weights: dict, allow_partial_loading: bool = False)
                 )
             # Split the pre-fused in_proj_qkv [Q|K|V] into its Q/K/V sub-tensors
             # BEFORE rank-blocking. Passing the concatenated [Q|K|V] as one
-            # component makes _rank_block cut it at contiguous (key+key+value)/tp
+            # component cuts it at contiguous (key+key+value)/tp
             # boundaries, which scrambles Q/K/V across ranks at tp>1 (rank r would
             # get [Q[..], K[..]] instead of [Q_r|K_r|V_r]) -> garbage GDN output.
             # Each of Q, K, V, and Z is sharded independently. Mirror the
@@ -402,8 +375,8 @@ def preprocess_weights(self, weights: dict, allow_partial_loading: bool = False)
                         f"{tuple(parts[projection].shape)}, expected {expected_ba_shape}"
                     )
             q_t, k_t, v_t = torch.split(qkv_t, [key_dim, key_dim, value_dim], dim=0)
-            qkvz = _rank_block([q_t, k_t, v_t, parts["z"][:]], tp_size)
-            ba = _rank_block([parts["b"][:], parts["a"][:]], tp_size)
+            qkvz = concatenate_weights_by_tp_rank([q_t, k_t, v_t, parts["z"][:]], tp_size)
+            ba = concatenate_weights_by_tp_rank([parts["b"][:], parts["a"][:]], tp_size)
             expected_qkvz_rows = 2 * key_dim + 2 * value_dim
             if qkvz.shape[0] != expected_qkvz_rows:
                 raise ValueError(

tensorrt_llm/_torch/models/modeling_utils.py

@@ -33,6 +33,7 @@
 from ..modules.rms_norm import RMSNorm
 from ..moe.fused_moe import MoE, VanillaMoE, is_moe_weight_owner
 from ..speculative import SpecMetadata
+from ..utils import split
 from ._arch_index import (MODEL_ARCH_TO_MODULE, SPEC_MODE_TO_MODULE,
                           is_builtin_zoo_module)
 
@@ -108,6 +109,44 @@ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
         return func(*args, **kwargs)
 
 
+def concatenate_weights_by_tp_rank(components: list[torch.Tensor],
+                                   tp_size: int) -> torch.Tensor:
+    """Concatenate projection components in rank-major row order.
+
+    Args:
+        components: Tensors shaped ``[rows_i, ...]`` with matching trailing
+            dimensions. Each physical row count must be divisible by
+            ``tp_size``. Packed weights remain packed; callers must ensure
+            logical channel boundaries align with physical row boundaries.
+        tp_size: Number of tensor-parallel ranks.
+
+    Returns:
+        A tensor ordered as ``[c0_r0 | c1_r0 | ... | c0_r1 | c1_r1 | ...]``.
+        A contiguous row split gives each rank its slice of every component.
+        With one rank, this is a plain concatenation along dimension zero.
+    """
+    if not components:
+        raise ValueError("rank blocking requires at least one projection")
+    if tp_size <= 0:
+        raise ValueError(f"tp_size must be positive, got {tp_size}")
+    trailing_shape = components[0].shape[1:]
+    for component in components:
+        if component.ndim == 0 or component.shape[1:] != trailing_shape:
+            raise ValueError(
+                "rank-blocked projections must have matching trailing dimensions"
+            )
+        if component.shape[0] % tp_size:
+            raise ValueError(
+                f"projection rows {component.shape[0]} are not divisible by tp_size={tp_size}"
+            )
+    if tp_size == 1:
+        return torch.cat(components, dim=0)
+    rows: list[torch.Tensor] = []
+    for rank in range(tp_size):
+        rows.extend(split(c, tp_size, rank) for c in components)
+    return torch.cat(rows, dim=0)
+
+
 def duplicate_kv_weight(weight: torch.Tensor, num_kv_heads: int,
                         tensor_parallel_size: int):
 

tensorrt_llm/_torch/modules/linear.py

@@ -2958,23 +2958,27 @@ def load_weights_vanilla(self, module: Linear, weights: List[Dict]) -> None:
         elm_packing = 2 if module.tp_mode == TensorParallelMode.COLUMN else 1
         load_weights_vanilla_helper(module, weights, elm_packing=elm_packing)
 
-        # Use the same device as the weight tensor
-        # as we register pre_quant_scale after sharded model weights are moved to respective gpus
-        device = module.weight.device
-        pre_quant_scale = module.load_shard(
-            weights[0]["pre_quant_scale"],
-            # pre_quant_scale applies to activation as opposed to weight, so flip tp_mode the other way around
-            flip_tp=True,
-            device=device,
-        )
+        pre_quant_scale = weights[0].get("pre_quant_scale")
+        if pre_quant_scale is not None:
+            # Use the device of the loaded weight for the activation scale.
+            device = module.weight.device
+            pre_quant_scale = module.load_shard(
+                pre_quant_scale,
+                # Activation scales shard along the input dimension.
+                flip_tp=True,
+                device=device,
+            )
 
-        module.pre_quant_scale = Parameter(
-            torch.ones((module.in_features, ), dtype=pre_quant_scale.dtype),
-            requires_grad=False).to(device=device)
+            module.pre_quant_scale = Parameter(
+                torch.ones((module.in_features, ), dtype=pre_quant_scale.dtype),
+                requires_grad=False).to(device=device)
+
+            copy_weight(module.pre_quant_scale, pre_quant_scale)
+        else:
+            module.pre_quant_scale = None
 
         weight_scale = self.load_weight_scales(module, weights)[0]
 
-        copy_weight(module.pre_quant_scale, pre_quant_scale)
         copy_weight(module.weight_scale, weight_scale.T.contiguous())
 
     def load_weights_fused_qkv_linear(self, module: Linear,

tests/integration/test_lists/test-db/l0_b200.yml

@@ -31,6 +31,8 @@ l0_b200:
       backend: pytorch
   tests:
   # ------------- PyTorch tests ---------------
+  - unittest/_torch/models/test_qwen3_5_int4_loading.py
+  - unittest/_torch/modules/test_linear_uneven_tp.py::TestW4A16AWQMLP::test_optional_pre_quant_scale
   - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_horizontal_producer.py
   - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_indexer.py
   - unittest/_torch/thop/parallel_hw_agnostic/test_minimax_m3_fp8_main_kv_insert.py