diff --git a/probe/issue-5/cutlass_ops_probe.py b/probe/issue-5/cutlass_ops_probe.py new file mode 100644 index 000000000..0b4781b8c --- /dev/null +++ b/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()) diff --git a/probe/issue-5/quantize_dit_int8wo.py b/probe/issue-5/quantize_dit_int8wo.py index 6b29f96f7..c929e40af 100644 --- a/probe/issue-5/quantize_dit_int8wo.py +++ b/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", }, diff --git a/probe/issue-5/w8a8_isolate.py b/probe/issue-5/w8a8_isolate.py new file mode 100644 index 000000000..53ab3d53c --- /dev/null +++ b/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())