probe(round two): W8A8 int8-dynamic scheme — kernel isolation, raw-op probe, artifact builder support

probe/issue-5/cutlass_ops_probe.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python
+"""Issue 5 round two: raw-op probe of the cutlass W8A8 int8 path on this stack.
+
+Bypasses every layer/loader and tests exactly what CutlassInt8ScaledMMLinearKernel
+calls: ops.scaled_int8_quant (dynamic per-token symmetric) then
+ops.cutlass_scaled_mm, against a float32 reference. A failure here means the
+cutlass int8 kernels themselves are broken on this stack (CUDA 13.4 / sm_86 /
+torch 2.13.0+cu130 / vllm 0.31.0), independent of any checkpoint or loader.
+
+Usage: cutlass_ops_probe.py
+"""
+
+from __future__ import annotations
+
+import json
+
+import torch
+from vllm import _custom_ops as ops
+
+
+def rel(a: torch.Tensor, b: torch.Tensor) -> float:
+    return float((a - b).norm() / b.norm())
+
+
+def main() -> int:
+    torch.manual_seed(2301)
+    device = torch.device("cuda")
+    results = []
+
+    for m, k, n in ((64, 256, 128), (512, 4096, 4096), (512, 4096, 12288)):
+        w = torch.randn(n, k, device=device)
+        w_scale = w.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
+        w_q = torch.round(w / w_scale).clamp(-127, 127).to(torch.int8)
+
+        x = torch.randn(m, k, dtype=torch.bfloat16, device=device)
+        x_q, x_s, _ = ops.scaled_int8_quant(x.contiguous(), None, None, symmetric=True)
+
+        # Reference: dequantized per-token acts @ dequantized per-channel weights.
+        ref = (x_q.float() * x_s.float()) @ (w_q.float() * w_scale.float()).T
+
+        w_q_t = w_q.t()
+        out = ops.cutlass_scaled_mm(
+            x_q, w_q_t, scale_a=x_s, scale_b=w_scale, out_dtype=torch.bfloat16
+        ).float()
+
+        # Also check the quantizer op itself against a float32 reference.
+        act_scale_ref = x.float().abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
+        x_q_ref = torch.round(x.float() / act_scale_ref).clamp(-127, 127).to(torch.int8)
+        quant_match = bool(torch.equal(x_q, x_q_ref)) and bool(
+            torch.allclose(x_s.float(), act_scale_ref, rtol=1e-3, atol=1e-7)
+        )
+
+        results.append({
+            "mk": [m, k, n],
+            "quant_op_matches_reference": quant_match,
+            "gemm_rel_l2": round(rel(out, ref), 5),
+            "out_finite": bool(torch.isfinite(out).all()),
+            "scale_a_shape": list(x_s.shape),
+            "scale_b_shape": list(w_scale.shape),
+        })
+
+    print(json.dumps(results, indent=1))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

probe/issue-5/quantize_dit_int8wo.py

@@ -61,8 +61,14 @@ SYMLINKED_COMPONENTS = ("scheduler", "text_encoder", "vae", "processor")
 SYMLINKED_FILES = ("model_index.json",)
 
 
-def build_quantization_dict() -> dict:
-    """Serialized compressed-tensors int8 weight-only metadata (llm-compressor schema)."""
+def build_quantization_dict(scheme: str) -> dict:
+    """Serialized compressed-tensors metadata (llm-compressor schema).
+
+    scheme="int8wo": int8 weight-only, pack-quantized storage.
+    scheme="w8a8dyn": int8 weights (static per-channel) + dynamic per-token
+        symmetric int8 activations, int-quantized storage — the non-Marlin
+        cutlass path on sm_86.
+    """
     weights = QuantizationArgs(
         num_bits=8,
         type=QuantizationType.INT,
@@ -71,10 +77,21 @@ def build_quantization_dict() -> dict:
         dynamic=False,
         observer="minmax",
     )
-    scheme = QuantizationScheme(targets=["Linear"], weights=weights)
+    input_activations = (
+        QuantizationArgs(
+            num_bits=8,
+            type=QuantizationType.INT,
+            symmetric=True,
+            strategy=QuantizationStrategy.TOKEN,
+            dynamic=True,
+        )
+        if scheme == "w8a8dyn"
+        else None
+    )
+    scheme_obj = QuantizationScheme(targets=["Linear"], weights=weights, input_activations=input_activations)
     config = QuantizationConfig(
-        config_groups={"group_0": scheme},
-        format=CompressionFormat.pack_quantized,
+        config_groups={"group_0": scheme_obj},
+        format=CompressionFormat.int_quantized if scheme == "w8a8dyn" else CompressionFormat.pack_quantized,
         ignore=list(IGNORED_TARGETS),
         quant_method="compressed-tensors",
         quantization_status=QuantizationStatus.FROZEN,
@@ -90,26 +107,39 @@ def sha256_file(path: Path) -> str:
     return digest.hexdigest()
 
 
-def quantize_weight(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, dict]:
-    """Per-output-channel symmetric int8: q = round(W/s), s = rowmax(|W|)/127.
+def quantize_weight(weight: torch.Tensor, scheme: str) -> tuple[dict[str, torch.Tensor], dict]:
+    """Per-output-channel symmetric int8.
 
-    Returns (weight_packed int32 [out, in/4], weight_scale fp32 [out, 1],
-    stats) using compressed-tensors' packer (uint8 biased by 128, 4 per int32).
+    scheme="int8wo": q = round(W/s), s = rowmax(|W|)/127, packed via
+    compressed-tensors' packer (uint8 biased by 128, 4 per int32).
+    scheme="w8a8dyn": same quantization, stored unpacked as plain int8
+    (int-quantized format) with fp32 scales.
     """
     w = weight.detach().to(torch.float32)
     scale = w.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
     q = torch.round(w / scale).clamp(-127, 127).to(torch.int8)
-    packed = pack_to_int32(q.contiguous(), 8).contiguous()
     dequant = q.to(torch.float32) * scale
     err = (w - dequant).norm() / w.norm().clamp_min(1e-12)
     stats = {"rel_l2": float(err), "out": int(w.shape[0]), "in": int(w.shape[1])}
-    return packed, scale.to(torch.float32).contiguous(), stats
+    if scheme == "w8a8dyn":
+        return {"weight": q.contiguous(), "weight_scale": scale.to(torch.float32).contiguous()}, stats
+    return {
+        "weight_packed": pack_to_int32(q.contiguous(), 8).contiguous(),
+        "weight_scale": scale.to(torch.float32).contiguous(),
+        "weight_shape": torch.tensor(weight.shape, dtype=torch.int64),
+    }, stats
 
 
 def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--snapshot", type=Path, required=True, help="sealed snapshot dir")
     parser.add_argument("--out", type=Path, required=True, help="artifact dir to create")
+    parser.add_argument(
+        "--scheme",
+        choices=("int8wo", "w8a8dyn"),
+        default="int8wo",
+        help="int8 weight-only (pack-quantized, Marlin path) or W8A8 int8-dynamic (int-quantized, cutlass path)",
+    )
     args = parser.parse_args()
 
     snapshot = args.snapshot.resolve()
@@ -131,10 +161,16 @@ def main() -> int:
         tensors = load_file(str(snapshot / "transformer" / shard))
         for name, tensor in tensors.items():
             if QUANTIZED_KEY.match(name):
-                packed, scale, stats = quantize_weight(tensor)
-                state[f"{name}_packed"] = packed
-                state[f"{name}_scale"] = scale
-                state[f"{name}_shape"] = torch.tensor(tensor.shape, dtype=torch.int64)
+                emitted, stats = quantize_weight(tensor, args.scheme)
+                if args.scheme == "w8a8dyn":
+                    # int-quantized storage: plain int8 weight + scale, mirroring
+                    # the vLLM layer parameter names (weight / weight_scale).
+                    state[name] = emitted["weight"]
+                    state[f"{name}_scale"] = emitted["weight_scale"]
+                else:
+                    state[f"{name}_packed"] = emitted["weight_packed"]
+                    state[f"{name}_scale"] = emitted["weight_scale"]
+                    state[f"{name}_shape"] = emitted["weight_shape"]
                 quant_stats[name] = stats
                 quant_params += tensor.numel()
                 print(f"q  {name} {tuple(tensor.shape)} rel_l2={stats['rel_l2']:.5f}", flush=True)
@@ -148,7 +184,7 @@ def main() -> int:
     total_bytes = weights_path.stat().st_size
 
     config = json.loads((snapshot / "transformer" / "config.json").read_text())
-    config["quantization_config"] = {"transformer": build_quantization_dict()}
+    config["quantization_config"] = {"transformer": build_quantization_dict(args.scheme)}
     (artifact / "transformer" / "config.json").write_text(json.dumps(config, indent=2) + "\n")
 
     for name in SYMLINKED_FILES:
@@ -166,9 +202,17 @@ def main() -> int:
         entry["sha256"] = sha256_file(artifact / entry["path"])
     tree_json = snapshot.parent.parent / "trees" / f"{SNAPSHOT_COMMIT}.json"
     manifest = {
-        "schema": "issue5-int8wo-v1",
+        "schema": f"issue5-{args.scheme}-v1",
         "created_utc": dt.datetime.now(dt.UTC).replace(microsecond=0).isoformat(),
-        "description": "Qwen-Image-2.1 DiT int8 weight-only (W8A16), compressed-tensors pack-quantized",
+        "description": (
+            f"Qwen-Image-2.1 DiT {args.scheme} "
+            + (
+                "W8A8 int8-dynamic, int-quantized"
+                if args.scheme == "w8a8dyn"
+                else "int8 weight-only W8A16, pack-quantized"
+            )
+            + ", compressed-tensors"
+        ),
         "inputs": {
             "snapshot": str(snapshot),
             "revision": SNAPSHOT_COMMIT,
@@ -184,7 +228,7 @@ def main() -> int:
             "method": "compressed-tensors",
             "format": "pack-quantized",
             "weights": "int8 symmetric per-output-channel, static, minmax",
-            "config": build_quantization_dict(),
+            "config": build_quantization_dict(args.scheme),
             "ignored_layers": list(IGNORED_TARGETS),
             "note": "ignore list mirrors the vllm-omni model's quant_config=None boundary layers",
         },

probe/issue-5/w8a8_isolate.py

@@ -0,0 +1,174 @@
+#!/usr/bin/env python
+"""Issue 5 round two: isolate the compressed-tensors W8A8-int8-dynamic path.
+
+Mirrors marlin_isolate.py but for the int-quantized W8A8 dynamic scheme
+(int8 weights, static per-output-channel scales, dynamic symmetric per-token
+int8 activations) whose kernel on sm_86 is the cutlass int8 scaled-MM —
+no Marlin in the path. Loads synthetic int8 tensors through the production
+weight-loader machinery and compares layer(x) against a float32 reference
+computed from the same dequantized weights with per-token dynamic int8
+activation quantization applied to the reference input.
+
+Expected rel-L2 when the chain is correct: a few percent (int8 weight error
+~1% plus int8 per-token activation rounding).
+
+Usage: w8a8_isolate.py [--block N]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+
+import torch
+
+
+def synth(out: int, inp: int, ones: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-channel symmetric int8 weight + scale, mirroring the artifact builder.
+
+    ones=True emits an analytically known weight (q=127, scale=1/127 -> W=1):
+    with an all-ones input row the layer output must be exactly in_features.
+    """
+    if ones:
+        q = torch.full((out, inp), 127, dtype=torch.int8)
+        return q, torch.full((out, 1), 1.0 / 127.0, dtype=torch.float32)
+    w = torch.randn(out, inp, dtype=torch.float32)
+    scale = w.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
+    q = torch.round(w / scale).clamp(-127, 127).to(torch.int8)
+    return q, scale.to(torch.float32)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--block", type=int, default=0)
+    parser.add_argument("--ones", action="store_true",
+                        help="analytic weights (W=1) and ones input; output must equal in_features")
+    parser.add_argument("--act-dtype", choices=("bfloat16", "float16"), default="bfloat16")
+    args = parser.parse_args()
+
+    import vllm.distributed
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.distributed import (
+        destroy_distributed_environment,
+        destroy_model_parallel,
+        initialize_model_parallel,
+    )
+    from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
+    from compressed_tensors.quantization import (
+        QuantizationArgs,
+        QuantizationConfig,
+        QuantizationScheme,
+        QuantizationStatus,
+        QuantizationStrategy,
+        QuantizationType,
+    )
+    from compressed_tensors.config import CompressionFormat
+    from vllm_omni.quantization import build_quant_config, resolve_component_quant_config
+
+    weights = QuantizationArgs(num_bits=8, type=QuantizationType.INT, symmetric=True,
+                               strategy=QuantizationStrategy.CHANNEL, dynamic=False, observer="minmax")
+    acts = QuantizationArgs(num_bits=8, type=QuantizationType.INT, symmetric=True,
+                            strategy=QuantizationStrategy.TOKEN, dynamic=True)
+    scheme = QuantizationScheme(targets=["Linear"], weights=weights, input_activations=acts)
+    disk = QuantizationConfig(
+        config_groups={"group_0": scheme}, format=CompressionFormat.int_quantized,
+        ignore=["img_in", "txt_in", "time_text_embed", "modulation", "norm_out", "proj_out"],
+        quant_method="compressed-tensors", quantization_status=QuantizationStatus.FROZEN,
+    ).to_dict()
+    quant_config = resolve_component_quant_config(
+        build_quant_config({"transformer": disk}), "transformer"
+    )
+
+    vllm.distributed.init_distributed_environment(
+        backend="gloo", world_size=1, rank=0, local_rank=0,
+        distributed_init_method="tcp://127.0.0.1:29602",
+    )
+    with set_current_vllm_config(VllmConfig()):
+        initialize_model_parallel(tensor_model_parallel_size=1)
+
+    device = torch.device("cuda")
+    torch.manual_seed(2301)
+    tensors = {
+        "q": synth(4096, 4096, args.ones),
+        "k": synth(4096, 4096, args.ones),
+        "v": synth(4096, 4096, args.ones),
+        "to_out": synth(4096, 4096, args.ones),
+    }
+
+    b = args.block
+    p = f"transformer.transformer_blocks.{b}"
+    results = []
+    with torch.device(device.type):
+        to_qkv = QKVParallelLinear(
+            hidden_size=4096, head_size=128, total_num_heads=32, bias=False,
+            quant_config=quant_config, prefix=f"{p}.attn.to_qkv",
+        )
+        for part, shard in (("q", "q"), ("k", "k"), ("v", "v")):
+            q, s = tensors[part]
+            to_qkv.weight.weight_loader(to_qkv.weight, q.to(device), shard)
+            to_qkv.weight_scale.weight_loader(to_qkv.weight_scale, s.to(device), shard)
+        to_qkv.quant_method.process_weights_after_loading(to_qkv)
+
+        to_out = RowParallelLinear(
+            4096, 4096, bias=False, input_is_parallel=True, return_bias=False,
+            quant_config=quant_config, prefix=f"{p}.attn.to_out.0",
+        )
+        q, s = tensors["to_out"]
+        to_out.weight.weight_loader(to_out.weight, q.to(device))
+        to_out.weight_scale.weight_loader(to_out.weight_scale, s.to(device))
+        loaded_bitwise_ok = bool(torch.equal(to_out.weight.detach().cpu(), q))
+        loaded_scale_ok = bool(torch.allclose(to_out.weight_scale.detach().float().cpu(), s))
+        to_out.quant_method.process_weights_after_loading(to_out)
+
+    for name, layer, parts, in_features in (
+        ("attn.to_qkv (fused q/k/v)", to_qkv, ("q", "k", "v"), 4096),
+        ("attn.to_out.0 (row)", to_out, ("to_out",), 4096),
+    ):
+        if args.ones:
+            x = torch.ones(4, in_features, dtype=getattr(torch, args.act_dtype), device=device)
+        else:
+            x = torch.randn(512, in_features, dtype=getattr(torch, args.act_dtype), device=device)
+        result = layer(x)
+        out = (result[0] if isinstance(result, tuple) else result).float()
+
+        if args.ones:
+            expected = torch.full_like(out, float(in_features))
+            rel = float((out - expected).norm() / expected.norm())
+            finite = bool(torch.isfinite(out).all())
+            results.append({
+                "layer": name,
+                "rel_l2": round(rel, 6) if finite else "NaN",
+                "out_finite": finite,
+                "out_row0_first8": [round(v, 2) for v in out[0, :8].tolist()],
+            })
+            continue
+
+        ref_w = torch.cat([tensors[k][0].float() * tensors[k][1] for k in parts], dim=0)
+        # Reference with the same dynamic per-token symmetric int8 activation
+        # quantization the kernel applies (dequantized acts, dequantized weights).
+        xs = x.float()
+        act_scale = xs.abs().amax(dim=1, keepdim=True).clamp_min(1e-12) / 127.0
+        x_dq = torch.round(xs / act_scale).clamp(-127, 127) * act_scale
+        ref = x_dq @ ref_w.T.to(device)
+        finite = bool(torch.isfinite(out).all())
+        rel = float((out - ref).norm() / ref.norm())
+        row = {"layer": name, "rel_l2": round(rel, 5) if finite else "NaN", "out_finite": finite}
+        if name.startswith("attn.to_out"):
+            row["loaded_bitwise_ok"] = loaded_bitwise_ok
+            row["loaded_scale_ok"] = loaded_scale_ok
+            # Reference computed from the layer's own post-repack weight
+            # parameter ((K, N) after the transpose) and scales.
+            w_layer = to_out.weight.detach().float()  # (K, N)
+            s_layer = to_out.weight_scale.detach().float().squeeze(1)  # (N,)
+            ref_layer = x_dq @ (w_layer * s_layer.unsqueeze(0)).to(device)
+            row["rel_l2_vs_layer_weight"] = round(float((out - ref_layer).norm() / ref_layer.norm()), 5)
+        results.append(row)
+
+    destroy_model_parallel()
+    destroy_distributed_environment()
+    print(json.dumps(results, indent=1))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())