[ROCm][Perf] W4A16: packed fp16 prefill dequant

benchmarks/kernels/benchmark_rdna_hybrid_w4a16_gemm.py

@@ -52,11 +52,11 @@
 # ---------------------------------------------------------------------------
 # Weight packing
 # ---------------------------------------------------------------------------
-def prepare_hybrid_weights(K, N, group_size, device="cuda"):
+def prepare_hybrid_weights(K, N, group_size, dtype=torch.float16, device="cuda"):
     """Create random weights for benchmarking.
 
-    Returns (w_q_skinny, w_s_skinny, w_fp16, w_zp). The triton path derives
-    its int32 view from w_q_skinny, so no separate int32 buffer is returned.
+    Returns (w_q_skinny, w_s_skinny, w_dense, w_zp). The triton path derives its
+    int32 view from w_q_skinny, so no separate int32 buffer is returned.
     """
     num_groups = K // group_size
 
@@ -65,23 +65,23 @@ def prepare_hybrid_weights(K, N, group_size, device="cuda"):
         0, 2**31, (N, K // 8), dtype=torch.int32, device=device
     )
     w_q_skinny = w_q_skinny_i32.view(torch.int8).contiguous()
-    w_s_skinny = torch.randn(N, num_groups, dtype=torch.float16, device=device) * 0.01
+    w_s_skinny = torch.randn(N, num_groups, dtype=dtype, device=device) * 0.01
 
     # Raw per-group zero-points for asymmetric benchmarks
     w_zp = torch.randint(0, 16, (N, num_groups), dtype=torch.int32, device=device).to(
-        torch.float16
+        dtype
     )
 
-    # FP16 baseline for F.linear
-    w_fp16 = torch.randn(N, K, dtype=torch.float16, device=device) * 0.01
+    # Unquantized baseline for F.linear
+    w_dense = torch.randn(N, K, dtype=dtype, device=device) * 0.01
 
-    return w_q_skinny, w_s_skinny, w_fp16, w_zp
+    return w_q_skinny, w_s_skinny, w_dense, w_zp
 
 
 # ---------------------------------------------------------------------------
 # Benchmark
 # ---------------------------------------------------------------------------
-PROVIDERS = ["torch-fp16", "hybrid-w4a16", "hybrid-w4a16-zp"]
+PROVIDERS = ["torch-dense", "hybrid-w4a16", "hybrid-w4a16-zp"]
 
 
 @triton.testing.perf_report(
@@ -93,22 +93,21 @@ def prepare_hybrid_weights(K, N, group_size, device="cuda"):
         line_vals=PROVIDERS,
         line_names=PROVIDERS,
         ylabel="TFLOP/s (larger is better)",
-        plot_name="FP16 vs Hybrid W4A16",
+        plot_name="Dense vs Hybrid W4A16",
         args={},
     )
 )
-def benchmark(batch_size, provider, N, K, group_size, weights):
+def benchmark(batch_size, provider, N, K, group_size, dtype, weights):
     M = batch_size
     device = "cuda"
-    dtype = torch.float16
     a = torch.randn((M, K), device=device, dtype=dtype)
 
     quantiles = [0.5, 0.2, 0.8]
 
-    if provider == "torch-fp16":
-        w_fp16 = weights["w_fp16"]
+    if provider == "torch-dense":
+        w_dense = weights["w_dense"]
         ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
-            lambda: torch.nn.functional.linear(a, w_fp16),
+            lambda: torch.nn.functional.linear(a, w_dense),
             quantiles=quantiles,
         )
     elif provider in ("hybrid-w4a16", "hybrid-w4a16-zp"):
@@ -168,21 +167,28 @@ def prepare_shapes(args):
     )
     parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
     parser.add_argument("--group-size", type=int, default=128)
+    parser.add_argument(
+        "--dtype", type=str, default="float16", choices=["float16", "bfloat16"]
+    )
     parser.add_argument("--save-path", type=str, default=None)
     args = parser.parse_args()
 
+    dtype = getattr(torch, args.dtype)
+
     for K, N, model in prepare_shapes(args):
         group_size = args.group_size
         print(f"\n{'=' * 70}")
-        print(f"{model}, N={N} K={K}, group_size={group_size}")
+        print(f"{model}, N={N} K={K}, group_size={group_size}, dtype={args.dtype}")
         print(f"{'=' * 70}")
 
-        w_q_skinny, w_s_skinny, w_fp16, w_zp = prepare_hybrid_weights(K, N, group_size)
+        w_q_skinny, w_s_skinny, w_dense, w_zp = prepare_hybrid_weights(
+            K, N, group_size, dtype
+        )
 
         weights = {
             "w_q_skinny": w_q_skinny,
             "w_s_skinny": w_s_skinny,
-            "w_fp16": w_fp16,
+            "w_dense": w_dense,
             "w_zp": w_zp,
         }
 
@@ -195,6 +201,7 @@ def prepare_shapes(args):
             N=N,
             K=K,
             group_size=group_size,
+            dtype=dtype,
             weights=weights,
         )
 

tests/kernels/quantization/test_rdna_hybrid_w4a16.py

@@ -30,6 +30,8 @@
 pack_int4_exllama_shuffle = hybrid_module.pack_int4_exllama_shuffle
 SUPPORTED_GROUP_SIZES = hybrid_module.SUPPORTED_GROUP_SIZES
 MAX_SKINNY_BATCH_SIZE = hybrid_module.MAX_SKINNY_BATCH_SIZE
+triton_w4a16_skinny_fmt_gemm = hybrid_module.triton_w4a16_skinny_fmt_gemm
+select_skinny_gfx1151_config = hybrid_module._select_skinny_gfx1151_config
 
 
 # ---------------------------------------------------------------------------
@@ -173,6 +175,78 @@ def test_rdna_hybrid_w4a16_apply_with_bias(dtype, M):
     torch.testing.assert_close(out, ref, rtol=2e-2, atol=2e-2)
 
 
+# ---------------------------------------------------------------------------
+# Triton prefill path
+# ---------------------------------------------------------------------------
+
+
+def _make_prefill_case(M, K, N, G, dtype, has_zp):
+    """Random [M,K] activations + skinny [N,K//8] weights and their metadata."""
+    x = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
+    w_int4 = torch.randint(0, 16, (N, K), device=device, dtype=torch.int32)
+    b_q = pack_int4_exllama_shuffle(w_int4)
+    scales = (0.05 * torch.rand((N, K // G), device=device, dtype=torch.float32)).to(
+        dtype
+    )
+    zp = (
+        torch.randint(0, 16, (N, K // G), device=device, dtype=torch.int32).to(dtype)
+        if has_zp
+        else None
+    )
+    return x, w_int4, b_q, scales, zp
+
+
+@pytest.mark.skipif(not on_gfx1x(), reason="Hybrid path is gfx11/gfx12 only")
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("has_zp", [False, True])
+@pytest.mark.parametrize(
+    "M,K,N,G",
+    [(17, 256, 512, 32), (32, 512, 256, 64), (33, 512, 512, 128), (64, 1024, 256, 128)],
+)
+def test_triton_prefill_gemm_matches_reference(dtype, has_zp, M, K, N, G):
+    """Prefill GEMM against a float32 oracle, over both unpacks and both the
+    asymmetric and symmetric dequants."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA/HIP device not available")
+    set_random_seed(0)
+
+    x, w_int4, b_q, scales, zp = _make_prefill_case(M, K, N, G, dtype, has_zp)
+    out = triton_w4a16_skinny_fmt_gemm(a=x, b_q=b_q, scales=scales, group_size=G, zp=zp)
+    ref = _rdna_hybrid_w4a16_reference(x, w_int4, scales, zp, G, bias=None)
+    torch.testing.assert_close(out, ref, rtol=1e-2, atol=5e-2)
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_gfx1151_tile_table_never_straddles_a_quant_group(dtype):
+    """BLOCK_K > group_size would give a tile's tail the wrong scale.
+
+    The kernel loads one scale per BLOCK_K tile, so this is a correctness
+    invariant of the table, not a tuning preference. Checked in Python so it
+    holds for shapes no test has hardware for.
+    """
+    for group_size in SUPPORTED_GROUP_SIZES:
+        for M in (1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096):
+            for N, K in [
+                (512, 2048),
+                (4096, 4096),
+                (24576, 4096),
+                (4096, 12288),
+                (32768, 2048),
+                (1024, 8192),
+            ]:
+                _, _, block_k, _, _ = select_skinny_gfx1151_config(
+                    M, N, K, group_size, dtype
+                )
+                assert block_k <= group_size, (
+                    f"BLOCK_K={block_k} > group_size={group_size} "
+                    f"at M={M} N={N} K={K} dtype={dtype}"
+                )
+                assert block_k % 8 == 0, (
+                    f"BLOCK_K={block_k} must be a multiple of 8 "
+                    f"(8 nibbles per packed int32)"
+                )
+
+
 # ---------------------------------------------------------------------------
 # pack_int4_exllama_shuffle round-trips correctly
 # ---------------------------------------------------------------------------