Remove redundant quant_storage parametrization from test_gemv_4bit

quant_storage is a dtype view over byte-identical packed 4-bit data,
introduced for FSDP sharding compatibility. The gemv kernel reads raw
bytes via void pointer and never branches on B.dtype, so testing 4
quant_storage values exercises the same code path 4 times.

Fix test_gemv_4bit to use only uint8 (the default), reducing from
1,536 to 384 test cases (~18 min saved). Add a dedicated
test_quant_storage_shard_roundtrip that properly tests what
quant_storage is for: verifying bytes survive FSDP-style
flatten/chunk/reassemble with all 4 storage dtypes.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

tests/test_functional.py

@@ -1281,13 +1281,9 @@ def test_bench_4bit_dequant(self, quant_type):
     @pytest.mark.parametrize("storage_type", ["nf4", "fp4"])
     @pytest.mark.parametrize("kind", ["fc1", "fc2", "attn", "attn_packed"])
     @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=describe_dtype)
-    @pytest.mark.parametrize(
-        "quant_storage",
-        [torch.uint8, torch.float16, torch.bfloat16, torch.float32],
-        ids=describe_dtype,
-    )
     @pytest.mark.parametrize("dim", [128, 256, 512, 1024], ids=id_formatter("dim"))
-    def test_gemv_4bit(self, device, dim, dtype, storage_type, quant_storage, double_quant, kind):
+    def test_gemv_4bit(self, device, dim, dtype, storage_type, double_quant, kind):
+        quant_storage = torch.uint8
         if device == "hpu" and not is_supported_on_hpu(storage_type, dtype, quant_storage):
             pytest.skip("This configuration is not supported on HPU.")
 

tests/test_linear4bit.py

@@ -248,6 +248,42 @@ def test_params4bit_torch_chunk_split(device, quant_type):
         assert split.quant_type == params4bit.quant_type, "Should preserve quant_type value"
 
 
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
+@pytest.mark.parametrize(
+    "quant_storage",
+    [torch.uint8, torch.float16, torch.bfloat16, torch.float32],
+    ids=describe_dtype,
+)
+def test_quant_storage_shard_roundtrip(device, quant_type, quant_storage):
+    """Test that quantized weights survive a flatten-chunk-reassemble roundtrip.
+
+    Non-uint8 quant_storage exists so that FSDP can shard quantized tensors
+    without splitting packed 4-bit pairs. This test simulates FSDP's
+    shard/gather pattern and verifies numerical correctness after reassembly.
+    """
+    M, K = 256, 128
+    A = torch.randn(1, K, dtype=torch.float16, device=device)
+    B = torch.randn(M, K, dtype=torch.float16, device=device)
+
+    qB, state = bnb.functional.quantize_4bit(B, quant_type=quant_type, quant_storage=quant_storage)
+    ref = bnb.functional.gemv_4bit(A, qB.t(), state=state)
+
+    # Simulate FSDP: flatten, split into shards, reassemble
+    flat = qB.flatten()
+    n_shards = 4
+    shards = flat.chunk(n_shards)
+    reassembled = torch.cat(shards).reshape(qB.shape)
+
+    assert reassembled.dtype == qB.dtype
+    assert torch.equal(
+        reassembled.view(torch.uint8), qB.view(torch.uint8)
+    ), "Bytes changed after shard roundtrip"
+
+    out = bnb.functional.gemv_4bit(A, reassembled.t(), state=state)
+    torch.testing.assert_close(out, ref)
+
+
 @pytest.mark.parametrize("device", get_available_devices())
 @pytest.mark.parametrize("quant_type", ["nf4", "fp4"])
 @pytest.mark.parametrize("blocksize", [32, 64, 128] if not ROCM_WARP_SIZE_64 else [64, 128])