[Bugfix][Loader] Read index shards in subfolders; direct I/O U16/U32

tests/model_executor/model_loader/test_direct_io.py

@@ -28,6 +28,13 @@ def _checkpoint(path) -> dict[str, torch.Tensor]:
         "model.language_model.layers.0.w": torch.randn(33, 17),
         "model.language_model.layers.1.w": torch.randn(64, 64).half(),
         "model.language_model.layers.2.w": torch.randn(5, 7).bfloat16(),
+        # csf_32x2_v1 page bases and offsets (GLM-5.3-Flash-FP4-Local)
+        "model.language_model.layers.1.page_bases": torch.arange(
+            -7, 9, dtype=torch.int32
+        ).view(torch.uint32),
+        "model.language_model.layers.2.offsets": torch.arange(
+            -5, 13, dtype=torch.int16
+        ).view(torch.uint16),
         "layers.3.ffn.w": torch.randn(9, 11).to(torch.float8_e4m3fn),
         "layers.3.ffn.scale": torch.randn(3, 4).to(torch.float8_e8m0fnu),
         "mtp.layers.0.w": torch.randint(0, 255, (129,), dtype=torch.uint8),

tests/model_executor/model_loader/test_safetensors_index_filter.py

@@ -151,3 +151,73 @@ def test_prepare_weights_skips_filtering_without_index(tmp_path) -> None:
     *_, indexed_weights_by_file = _prepare(tmp_path, load_format="safetensors")
 
     assert indexed_weights_by_file is None
+
+
+def _write_subfolder_checkpoint(tmp_path) -> None:
+    """GLM-5.3-Flash-FP4-Local layout: one indexed shard lives in a subfolder."""
+    (tmp_path / "tensors").mkdir()
+    save_file({"main.weight": torch.tensor([1.0])}, tmp_path / "model.safetensors")
+    native = tmp_path / "tensors" / "native.safetensors"
+    save_file({"native.weight": torch.tensor([2.0])}, native)
+    _write_index(
+        tmp_path,
+        {
+            "main.weight": "model.safetensors",
+            "native.weight": "tensors/native.safetensors",
+        },
+    )
+
+
+@pytest.mark.parametrize("load_strategy", [None, "lazy", "eager"])
+def test_default_loader_reads_indexed_shards_in_subfolders(
+    tmp_path, load_strategy: str | None
+) -> None:
+    # A top-level glob misses the subfolder shard; without it the tensors it
+    # holds would stay unloaded without any error.
+    _write_subfolder_checkpoint(tmp_path)
+    loader = DefaultModelLoader(
+        LoadConfig(
+            load_format="safetensors",
+            safetensors_load_strategy=load_strategy,
+            use_tqdm_on_load=False,
+        )
+    )
+    source = DefaultModelLoader.Source(
+        model_or_path=str(tmp_path), revision=None, fall_back_to_pt=False
+    )
+
+    weights = dict(loader._get_weights_iterator(source))
+
+    assert sorted(weights) == ["main.weight", "native.weight"]
+    torch.testing.assert_close(weights["native.weight"], torch.tensor([2.0]))
+
+
+def test_prepare_weights_fails_when_an_indexed_shard_is_missing(tmp_path) -> None:
+    save_file({"main.weight": torch.tensor([1.0])}, tmp_path / "model.safetensors")
+    _write_index(
+        tmp_path,
+        {"main.weight": "model.safetensors", "gone.weight": "model-00002.safetensors"},
+    )
+
+    with pytest.raises(RuntimeError, match="would not be loaded"):
+        _prepare(tmp_path, load_format="safetensors")
+
+
+def test_prepare_weights_lists_each_subfolder_shard_once(tmp_path) -> None:
+    # FP4-Local layout: every indexed shard lives in a subfolder, so the
+    # index-only fallback finds them; the subfolder pass must not add them
+    # a second time.
+    (tmp_path / "tensors").mkdir()
+    weight_map = {}
+    for i in range(2):
+        shard = f"tensors/model-0000{i + 1}-of-00002.safetensors"
+        save_file({f"w{i}": torch.tensor([float(i)])}, tmp_path / shard)
+        weight_map[f"w{i}"] = shard
+    _write_index(tmp_path, weight_map)
+
+    _, files, use_safetensors, _ = _prepare(tmp_path, load_format="safetensors")
+
+    assert use_safetensors
+    assert sorted(os.path.relpath(f, tmp_path) for f in files) == sorted(
+        weight_map.values()
+    )

vllm/model_executor/model_loader/default_loader.py

@@ -34,6 +34,8 @@
     filter_files_not_needed_for_inference,
     get_quant_config,
     get_safetensors_index_weights_by_file,
+    indexed_shards_in_subfolders,
+    indexed_shards_not_selected,
     instanttensor_weights_iterator,
     maybe_download_from_modelscope,
     multi_thread_pt_weights_iterator,
@@ -260,6 +262,16 @@ def _prepare_weights(
             if len(hf_weights_files) > 0:
                 if pattern.endswith(".safetensors"):
                     use_safetensors = True
+                    # The index-only fallback above already lists every
+                    # shard; add only the subfolder shards not yet selected.
+                    selected = {os.path.normpath(f) for f in hf_weights_files}
+                    hf_weights_files += [
+                        shard
+                        for shard in indexed_shards_in_subfolders(
+                            hf_folder, index_file, pattern
+                        )
+                        if os.path.normpath(shard) not in selected
+                    ]
                 break
 
         if use_safetensors:
@@ -279,6 +291,19 @@ def _prepare_weights(
             hf_weights_files = filter_duplicate_safetensors_files(
                 hf_weights_files, hf_folder, index_file
             )
+            # A shard the index names but the loader would skip leaves its
+            # parameters unloaded without any error; only a model's own
+            # allow patterns may leave shards out on purpose.
+            unread = (
+                indexed_shards_not_selected(hf_folder, index_file, hf_weights_files)
+                if allow_patterns_overrides is None
+                else []
+            )
+            if unread:
+                raise RuntimeError(
+                    f"The checkpoint index of {model_name_or_path} names shards "
+                    f"that would not be loaded: {unread[:4]}"
+                )
         else:
             hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files)
 

vllm/model_executor/model_loader/direct_io.py

@@ -52,6 +52,10 @@
     "I16": torch.int16,
     "I8": torch.int8,
     "U8": torch.uint8,
+    # Unsigned index tables, e.g. GLM-5.3's csf_32x2_v1 scale pages; mapped
+    # like safetensors maps them.
+    "U16": torch.uint16,
+    "U32": torch.uint32,
     "BOOL": torch.bool,
 }
 

vllm/model_executor/model_loader/weight_utils.py

@@ -689,6 +689,48 @@ def filter_duplicate_safetensors_files(
     return hf_weights_files
 
 
+def _index_shards(hf_folder: str, index_file: str) -> set[str] | None:
+    index_file_name = os.path.join(hf_folder, index_file)
+    if not os.path.isfile(index_file_name):
+        return None
+    with open(index_file_name) as f:
+        return set(json.load(f)["weight_map"].values())
+
+
+def indexed_shards_in_subfolders(
+    hf_folder: str, index_file: str, pattern: str
+) -> list[str]:
+    """Shards the checkpoint index keeps in subfolders whose file name matches
+    a top-level ``pattern``: the folder glob does not reach them.
+
+    (GLM-5.3-Flash-FP4-Local stores its BF16 attention tensors in
+    ``tensors/native-QB.safetensors``.)
+    """
+    shards = _index_shards(hf_folder, index_file)
+    if shards is None or "/" in pattern:
+        return []
+    return sorted(
+        os.path.join(hf_folder, shard)
+        for shard in shards
+        if os.path.dirname(shard) and fnmatch.fnmatch(os.path.basename(shard), pattern)
+    )
+
+
+def indexed_shards_not_selected(
+    hf_folder: str, index_file: str, hf_weights_files: list[str]
+) -> list[str]:
+    """Shards the checkpoint index names that are not among the files to load."""
+    shards = _index_shards(hf_folder, index_file)
+    if shards is None:
+        return []
+    selected = {os.path.normpath(path) for path in hf_weights_files}
+    return sorted(
+        shard
+        for shard in shards
+        if os.path.normpath(os.path.join(hf_folder, shard)) not in selected
+    )
+
+
 def get_safetensors_index_weights_by_file(
     hf_folder: str, index_file: str, hf_weights_files: list[str]
 ) -> dict[str, set[str]] | None: