From d2a742f711a7390fd368e9ea39e37104c43e2ac6 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 8 Jul 2026 16:09:21 -0400 Subject: [PATCH 1/7] Basic convrot int4 support. Should work on 30 and 40 series. 50 series doesn't have int4 ops so it will use int8 ops. 20 series needs special code for int4 and should fallback to int8 at the moment but is untested. --- comfy_kitchen/__init__.py | 8 + comfy_kitchen/backends/cuda/CMakeLists.txt | 1 + comfy_kitchen/backends/cuda/__init__.py | 464 +++ .../backends/cuda/dlpack_bindings.cpp | 409 +++ .../backends/cuda/ops/convrot_w4a4.cu | 2650 +++++++++++++++++ .../backends/cuda/ops/cutlass_gemm_int8.cu | 2 +- comfy_kitchen/backends/eager/__init__.py | 46 + comfy_kitchen/backends/eager/convrot_w4a4.py | 196 ++ comfy_kitchen/tensor/__init__.py | 11 + comfy_kitchen/tensor/convrot_w4a4.py | 244 ++ 10 files changed, 4030 insertions(+), 1 deletion(-) create mode 100644 comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu create mode 100644 comfy_kitchen/backends/eager/convrot_w4a4.py create mode 100644 comfy_kitchen/tensor/convrot_w4a4.py diff --git a/comfy_kitchen/__init__.py b/comfy_kitchen/__init__.py index a8172ef..96514d9 100644 --- a/comfy_kitchen/__init__.py +++ b/comfy_kitchen/__init__.py @@ -17,6 +17,11 @@ # Import registry and exceptions from .registry import registry +from .tensor.convrot_w4a4 import ( + convrot_w4a4_linear, + dequantize_convrot_w4a4_weight, + quantize_convrot_w4a4_weight, +) __version__ = "0.1.0" @@ -31,6 +36,7 @@ "quantize_mxfp8", "dequantize_mxfp8", "quantize_svdquant_w4a4", + "quantize_convrot_w4a4_weight", "quantize_int8_rowwise", "quantize_int8_tensorwise", "dequantize_int8_simple", @@ -38,6 +44,8 @@ "scaled_mm_nvfp4", "scaled_mm_mxfp8", "scaled_mm_svdquant_w4a4", + "convrot_w4a4_linear", + "dequantize_convrot_w4a4_weight", "gemv_awq_w4a16", "int8_linear", # Positional encoding diff --git a/comfy_kitchen/backends/cuda/CMakeLists.txt b/comfy_kitchen/backends/cuda/CMakeLists.txt index 9f2b1db..8beae7a 100644 --- a/comfy_kitchen/backends/cuda/CMakeLists.txt +++ b/comfy_kitchen/backends/cuda/CMakeLists.txt @@ -108,6 +108,7 @@ set(CUDA_SOURCES ops/cublas_gemm_int8.cu ops/cutlass_gemm_int8.cu ops/int8_linear.cu + ops/convrot_w4a4.cu ops/apply_rope.cu ops/quantize_svdquant_w4a4.cu ops/scaled_mm_svdquant_w4a4.cu diff --git a/comfy_kitchen/backends/cuda/__init__.py b/comfy_kitchen/backends/cuda/__init__.py index 228ec9c..5780dbf 100644 --- a/comfy_kitchen/backends/cuda/__init__.py +++ b/comfy_kitchen/backends/cuda/__init__.py @@ -32,9 +32,17 @@ "dequantize_int8_simple_dtype", "dequantize_int8_convrot_weight", "dequantize_int8_convrot_weight_dtype", + "dequantize_convrot_w4a4_weight", "int8_linear", + "int4_linear", + "convrot_w4a4_linear", + "prepare_int4_weight_for_int8_linear", "quantize_int8_tensorwise", "quantize_int8_rowwise", + "quantize_int4_rowwise", + "quantize_int4_rowwise_convrot64", + "quantize_int4_rowwise_convrot64_to_int8", + "quantize_convrot_w4a4_weight", "quantize_int8_convrot_weight", "quantize_int8_rowwise_convrot64", "quantize_and_rotate_rowwise", @@ -127,6 +135,10 @@ def find_lib_dir(start_dir, lib_pattern): from comfy_kitchen.backends.eager.quantization import ( # noqa: E402 quantize_int8_tensorwise as eager_quantize_int8_tensorwise, ) +from comfy_kitchen.backends.eager.svdquant import ( # noqa: E402 + _INT4_GROUP_SIZE, + _unpack_int4_row_major, +) from comfy_kitchen.constraints import ( # noqa: E402 DivisibleBy, ExactDims, @@ -148,6 +160,7 @@ def find_lib_dir(start_dir, lib_pattern): _turing_device_cache: dict[int, bool] = {} _cutlass_int8_device_cache: dict[int, bool] = {} _shared_memory_per_block_cache: dict[int, int] = {} +_FORCE_INT4_INT8_FALLBACK = os.environ.get("COMFY_KITCHEN_FORCE_INT4_INT8_FALLBACK", "0") == "1" def _cuda_device_is_turing(device_index: int) -> bool: @@ -172,6 +185,16 @@ def _cuda_device_supports_cutlass_int8_dequant(tensor: torch.Tensor) -> bool: return supported +def _cuda_device_supports_native_int4_mma(tensor: torch.Tensor) -> bool: + if not tensor.is_cuda or _FORCE_INT4_INT8_FALLBACK: + return False + major, _minor = torch.cuda.get_device_capability(tensor.get_device()) + # The current ConvRot W4A4 kernel emits m16n8k64 s4 MMA, which is the + # sm80+ integer MMA shape. Hopper is routed through the INT8 fallback for + # better behavior with this implementation. + return major == 8 + + def _cublas_int8_n_alignment(tensor: torch.Tensor) -> int: # Turing cuBLASLt INT8 rejects some skinny-N shapes, e.g. N=17. return 32 if tensor.is_cuda and _cuda_device_is_turing(tensor.get_device()) else 8 @@ -487,6 +510,387 @@ def quantize_int8_rowwise( return q_2d.reshape(orig_shape), scales_2d.reshape(*orig_shape[:-1], 1) +def quantize_int4_rowwise( + x: torch.Tensor, + stochastic_rounding: int | None = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a contiguous 2D tensor to signed int4 with one float scale per row.""" + orig_shape = x.shape + x_2d = x.reshape(-1, x.shape[-1]).contiguous() + if x_2d.shape[-1] % 64 != 0: + raise ValueError(f"INT4 rowwise quantization requires K divisible by 64, got {x_2d.shape[-1]}") + q_2d = torch.empty((x_2d.shape[0], x_2d.shape[1] // 2), dtype=torch.int8, device=x.device) + scales_2d = torch.empty((x_2d.shape[0], 1), dtype=torch.float32, device=x.device) + stream_ptr = torch.cuda.current_stream(x.device).cuda_stream + _C.quantize_int4_rowwise( + _wrap_for_dlpack(x_2d), + _wrap_for_dlpack(q_2d), + _wrap_for_dlpack(scales_2d), + stochastic_rounding is not None and stochastic_rounding > 0, + int(stochastic_rounding or 0), + stream_ptr, + ) + return q_2d.reshape(*orig_shape[:-1], orig_shape[-1] // 2), scales_2d.reshape(*orig_shape[:-1], 1) + + +def quantize_int4_rowwise_convrot64( + x: torch.Tensor, + group_size: int, + stochastic_rounding: int | None = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused regular ConvRot rotation plus rowwise signed int4 quantization.""" + orig_shape = x.shape + x_2d = x.reshape(-1, x.shape[-1]).contiguous() + if group_size not in (16, 64, 256): + raise ValueError(f"INT4 ConvRot fused quantization requires group_size 16, 64, or 256, got {group_size}") + if x_2d.shape[-1] % group_size != 0: + raise ValueError(f"INT4 ConvRot fused quantization requires K divisible by group_size {group_size}, got {x_2d.shape[-1]}") + q_2d = torch.empty((x_2d.shape[0], x_2d.shape[1] // 2), dtype=torch.int8, device=x.device) + scales_2d = torch.empty((x_2d.shape[0], 1), dtype=torch.float32, device=x.device) + stream_ptr = torch.cuda.current_stream(x.device).cuda_stream + _C.quantize_int4_rowwise_convrot64( + _wrap_for_dlpack(x_2d), + _wrap_for_dlpack(q_2d), + _wrap_for_dlpack(scales_2d), + group_size, + stochastic_rounding is not None and stochastic_rounding > 0, + int(stochastic_rounding or 0), + stream_ptr, + ) + return q_2d.reshape(*orig_shape[:-1], orig_shape[-1] // 2), scales_2d.reshape(*orig_shape[:-1], 1) + + +def quantize_int4_rowwise_convrot64_to_int8( + x: torch.Tensor, + group_size: int, + stochastic_rounding: int | None = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused ConvRot-256 rotation plus int4-scale quantization into INT8 storage.""" + orig_shape = x.shape + x_2d = x.reshape(-1, x.shape[-1]).contiguous() + if group_size != 256: + raise ValueError(f"INT4 fallback INT8 activation quantization requires group_size 256, got {group_size}") + if x_2d.shape[-1] % group_size != 0: + raise ValueError(f"INT4 ConvRot fallback quantization requires K divisible by {group_size}, got {x_2d.shape[-1]}") + q_2d = torch.empty_like(x_2d, dtype=torch.int8) + scales_2d = torch.empty((x_2d.shape[0], 1), dtype=torch.float32, device=x.device) + stream_ptr = torch.cuda.current_stream(x.device).cuda_stream + _C.quantize_int4_rowwise_convrot64_to_int8( + _wrap_for_dlpack(x_2d), + _wrap_for_dlpack(q_2d), + _wrap_for_dlpack(scales_2d), + group_size, + stochastic_rounding is not None and stochastic_rounding > 0, + int(stochastic_rounding or 0), + stream_ptr, + ) + return q_2d.reshape(*orig_shape), scales_2d.reshape(*orig_shape[:-1], 1) + + +def _unpack_int4_to_int8_cuda(qdata: torch.Tensor) -> torch.Tensor: + qdata_2d = qdata.contiguous() + output = torch.empty((qdata_2d.shape[0], qdata_2d.shape[1] * 2), dtype=torch.int8, device=qdata_2d.device) + stream_ptr = torch.cuda.current_stream(qdata_2d.device).cuda_stream + _C.unpack_int4_to_int8(_wrap_for_dlpack(qdata_2d), _wrap_for_dlpack(output), stream_ptr) + return output + + +def _int4_linear_via_int8_values( + x_int8: torch.Tensor, + weight_int8: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + if x_int8.dim() != 2 or weight_int8.dim() != 2: + raise ValueError("INT4 fallback INT8 GEMM expects 2D activation and weight tensors") + if x_int8.shape[1] != weight_int8.shape[1]: + raise ValueError("INT4 fallback INT8 GEMM K dimensions do not match") + m, k = x_int8.shape + n = weight_int8.shape[0] + output = torch.empty((m, n), dtype=out_dtype, device=x_int8.device) + + x_scale_arg = x_scale.to(device=x_int8.device, dtype=torch.float32).reshape(-1).contiguous() + weight_scale_arg = weight_scale.to(device=x_int8.device, dtype=torch.float32).reshape(-1).contiguous() + if x_scale_arg.numel() != m: + raise ValueError(f"INT4 fallback x_scale must have {m} values, got {x_scale_arg.numel()}") + if weight_scale_arg.numel() != n: + raise ValueError(f"INT4 fallback weight_scale must have {n} values, got {weight_scale_arg.numel()}") + bias_arg = bias if bias is not None else _empty_cuda_tensor(x_int8.device, out_dtype) + if bias is not None and (bias.device != x_int8.device or bias.dtype != out_dtype or not bias.is_contiguous()): + bias_arg = bias.to(device=x_int8.device, dtype=out_dtype).contiguous() + + used_cutlass = False + if ( + not _DISABLE_CUTLASS_INT8 + and _cuda_device_supports_cutlass_int8_dequant(x_int8) + and m <= 512 + and n <= k + and k <= 4096 + ): + ws_cutlass = weight_scale_arg if weight_scale_arg.numel() == n else weight_scale_arg.expand(n).contiguous() + bias_f32 = bias_arg.to(torch.float32).contiguous() if bias is not None else bias_arg + used_cutlass = _C.cutlass_int8_dequant( + _wrap_for_dlpack(x_int8), + _wrap_for_dlpack(weight_int8), + _wrap_for_dlpack(x_scale_arg.reshape(m, 1)), + _wrap_for_dlpack(ws_cutlass), + _wrap_for_dlpack(bias_f32), + _wrap_for_dlpack(output), + DTYPE_TO_CODE[out_dtype], + torch.cuda.current_stream(x_int8.device).cuda_stream, + ) + if used_cutlass: + return output + + if _cuda_device_is_turing(x_int8.get_device()): + padded_k = _round_up(k, 16) + padded_n = _round_up(n, _cublas_int8_n_alignment(x_int8)) + cublas_x = _pad_2d_cols(x_int8, padded_k) + cublas_weight = _pad_2d_rows(_pad_2d_cols(weight_int8, padded_k), padded_n) + else: + padded_n = n + cublas_x = x_int8 + cublas_weight = weight_int8 + + out_int32 = torch.empty((m, padded_n), dtype=torch.int32, device=x_int8.device) + stream_ptr = torch.cuda.current_stream(x_int8.device).cuda_stream + _C.cublas_gemm_int8( + _wrap_for_dlpack(cublas_x), + _wrap_for_dlpack(cublas_weight), + _wrap_for_dlpack(out_int32), + _wrap_for_dlpack(get_cublas_workspace()), + stream_ptr, + ) + if padded_n != n: + out_int32 = out_int32[:, :n].contiguous() + _C.dequantize_int8_linear( + out_int32, + x_scale_arg.reshape(m, 1), + weight_scale_arg, + bias_arg, + output, + DTYPE_TO_CODE[out_dtype], + stream_ptr, + ) + return output + + +def _int4_linear_via_int8( + x_qdata: torch.Tensor, + weight: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, + weight_int8: torch.Tensor | None = None, +) -> torch.Tensor: + _m, k_half = x_qdata.shape + n = weight.shape[0] + k = k_half * 2 + x_int8 = _unpack_int4_to_int8_cuda(x_qdata) + if weight_int8 is None: + weight_int8 = _unpack_int4_to_int8_cuda(weight) + elif weight_int8.shape != (n, k) or weight_int8.dtype != torch.int8 or weight_int8.device != weight.device: + raise ValueError("prepared INT8 fallback weight has incompatible shape, dtype, or device") + return _int4_linear_via_int8_values(x_int8, weight_int8, x_scale, weight_scale, bias, out_dtype) + + +def int4_linear( + x_qdata: torch.Tensor, + weight: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None = None, + out_dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + """Signed INT4 linear: output = (x_q @ weight.T) * x_scale * weight_scale + bias.""" + if x_qdata.dim() != 2 or weight.dim() != 2: + raise ValueError("INT4 linear expects 2D activation and weight tensors") + if x_qdata.shape[1] != weight.shape[1]: + raise ValueError("INT4 linear activation/weight K dimensions do not match") + m, _k_half = x_qdata.shape + n = weight.shape[0] + output = torch.empty((m, n), dtype=out_dtype, device=x_qdata.device) + x_scale_arg = x_scale.to(device=x_qdata.device, dtype=torch.float32).reshape(-1).contiguous() + weight_scale_arg = weight_scale.to(device=x_qdata.device, dtype=torch.float32).reshape(-1).contiguous() + if x_scale_arg.numel() != m: + raise ValueError(f"INT4 x_scale must have {m} values, got {x_scale_arg.numel()}") + if weight_scale_arg.numel() != n: + raise ValueError(f"INT4 weight_scale must have {n} values, got {weight_scale_arg.numel()}") + bias_arg = bias if bias is not None else _empty_cuda_tensor(x_qdata.device, out_dtype) + if bias is not None and (bias.device != x_qdata.device or not bias.is_contiguous()): + bias_arg = bias.to(device=x_qdata.device).contiguous() + stream_ptr = torch.cuda.current_stream(x_qdata.device).cuda_stream + if not _cuda_device_supports_native_int4_mma(x_qdata): + return _int4_linear_via_int8( + x_qdata.contiguous(), + weight.contiguous(), + x_scale_arg, + weight_scale_arg, + bias_arg if bias is not None else None, + out_dtype, + ) + used_cutlass = False + if ( + not _DISABLE_CUTLASS_INT8 + and _cuda_device_supports_native_int4_mma(x_qdata) + and _cuda_device_supports_cutlass_int8_dequant(x_qdata) + and hasattr(_C, "cutlass_int4_dequant") + ): + bias_f32 = bias_arg.to(torch.float32).contiguous() if bias is not None else bias_arg + used_cutlass = _C.cutlass_int4_dequant( + _wrap_for_dlpack(x_qdata.contiguous()), + _wrap_for_dlpack(weight.contiguous()), + _wrap_for_dlpack(x_scale_arg), + _wrap_for_dlpack(weight_scale_arg), + _wrap_for_dlpack(bias_f32), + _wrap_for_dlpack(output), + DTYPE_TO_CODE[out_dtype], + stream_ptr, + ) + if used_cutlass: + return output + _C.int4_linear( + _wrap_for_dlpack(x_qdata.contiguous()), + _wrap_for_dlpack(weight.contiguous()), + _wrap_for_dlpack(x_scale_arg), + _wrap_for_dlpack(weight_scale_arg), + _wrap_for_dlpack(bias_arg), + _wrap_for_dlpack(output), + DTYPE_TO_CODE[out_dtype], + stream_ptr, + ) + return output + + +def prepare_int4_weight_for_int8_linear(weight: torch.Tensor) -> torch.Tensor: + """Prepare packed signed INT4 weight as INT8 for non-native INT4 GEMM fallback.""" + if weight.dim() != 2: + raise ValueError("prepared INT4 fallback weight expects a 2D tensor") + return _unpack_int4_to_int8_cuda(weight) + + +def quantize_convrot_w4a4_weight( + weight: torch.Tensor, + convrot_groupsize: int = 256, + quant_group_size: int = _INT4_GROUP_SIZE, + stochastic_rounding: int | None = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Offline ConvRot weight rotation followed by row-wise signed INT4 quantization.""" + if quant_group_size != _INT4_GROUP_SIZE: + raise ValueError(f"int4 MMA kernel requires quant_group_size {_INT4_GROUP_SIZE}") + if weight.dim() != 2: + raise ValueError(f"ConvRot W4A4 expects a 2D tensor, got shape {tuple(weight.shape)}") + if weight.shape[-1] % convrot_groupsize != 0: + raise ValueError(f"in_features {weight.shape[-1]} not divisible by convrot_groupsize {convrot_groupsize}") + weight_2d = weight.contiguous() + if convrot_groupsize in (16, 64, 256) and hasattr(_C, "quantize_int4_rowwise_convrot64"): + qdata, scales = quantize_int4_rowwise_convrot64( + weight_2d, + convrot_groupsize, + stochastic_rounding=stochastic_rounding, + ) + else: + h = _build_hadamard(convrot_groupsize, device=weight.device, dtype=weight.dtype) + weight_rot = _rotate_weight(weight_2d, h, convrot_groupsize) + qdata, scales = quantize_int4_rowwise(weight_rot.contiguous(), stochastic_rounding=stochastic_rounding) + return qdata, scales.reshape(-1) + + +def dequantize_convrot_w4a4_weight( + qdata: torch.Tensor, + scales: torch.Tensor, + convrot_groupsize: int = 256, + quant_group_size: int = _INT4_GROUP_SIZE, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Dequantize packed ConvRot W4A4 weights and rotate back to original basis.""" + if quant_group_size != _INT4_GROUP_SIZE: + raise ValueError(f"int4 MMA kernel requires quant_group_size {_INT4_GROUP_SIZE}") + if qdata.dim() != 2: + raise ValueError(f"ConvRot W4A4 dequant expects 2D qdata, got shape {tuple(qdata.shape)}") + k = qdata.shape[-1] * 2 + if k % convrot_groupsize != 0: + raise ValueError(f"in_features {k} not divisible by convrot_groupsize {convrot_groupsize}") + if convrot_groupsize in (16, 64, 256) and hasattr(_C, "dequantize_int4_convrot64"): + qdata_2d = qdata.contiguous() + scale_arg = scales.to(device=qdata.device, dtype=torch.float32).reshape(-1).contiguous() + output = torch.empty((qdata_2d.shape[0], k), dtype=output_dtype, device=qdata.device) + stream_ptr = torch.cuda.current_stream(qdata.device).cuda_stream + _C.dequantize_int4_convrot64( + _wrap_for_dlpack(qdata_2d), + _wrap_for_dlpack(scale_arg), + _wrap_for_dlpack(output), + convrot_groupsize, + stream_ptr, + ) + return output + w_int = _unpack_int4_row_major(qdata).to(torch.float32) + w_rot = w_int * scales.to(device=qdata.device, dtype=torch.float32).reshape(-1, 1) + h = _build_hadamard(convrot_groupsize, device=qdata.device, dtype=torch.float32) + return _rotate_weight(w_rot.float(), h, convrot_groupsize).to(output_dtype) + + +def convrot_w4a4_linear( + x: torch.Tensor, + qweight: torch.Tensor, + wscales: torch.Tensor, + bias: torch.Tensor | None = None, + convrot_groupsize: int = 256, + quant_group_size: int = _INT4_GROUP_SIZE, +) -> torch.Tensor: + """Compute ``x @ W.T + bias`` using ConvRot W4A4 signed INT4 MMA.""" + if quant_group_size != _INT4_GROUP_SIZE: + raise ValueError(f"int4 MMA kernel requires quant_group_size {_INT4_GROUP_SIZE}") + if x.shape[-1] != qweight.shape[-1] * 2: + raise ValueError(f"Input K={x.shape[-1]} does not match qweight K={qweight.shape[-1] * 2}") + if x.shape[-1] % convrot_groupsize != 0: + raise ValueError(f"Input K={x.shape[-1]} not divisible by convrot_groupsize {convrot_groupsize}") + + orig_shape = x.shape + x2d = x.reshape(-1, orig_shape[-1]).contiguous() + if not _cuda_device_supports_native_int4_mma(x2d): + weight_int8 = prepare_int4_weight_for_int8_linear(qweight) + if ( + convrot_groupsize == 256 + and x2d.shape[-1] % 256 == 0 + and 256 <= x2d.shape[-1] <= _CONVROT_FUSED_MAX_K + and _convrot_fused_shared_memory_fits(x2d, x2d.shape[-1], convrot_groupsize) + ): + qact_int8, x_scale = quantize_int8_rowwise_convrot64(x2d, convrot_groupsize) + elif _should_use_convrot_fused_kernel(x2d, x2d.shape[-1], convrot_groupsize): + qact_int8, x_scale = quantize_int8_rowwise_convrot(x2d, convrot_groupsize) + else: + h = _build_hadamard(convrot_groupsize, device=x2d.device, dtype=x2d.dtype) + qact_int8, x_scale = quantize_and_rotate_rowwise(x2d, h, convrot_groupsize) + out = _int4_linear_via_int8_values( + qact_int8, + weight_int8, + x_scale, + wscales, + bias, + x.dtype, + ) + return out[: x2d.shape[0]].reshape(*orig_shape[:-1], qweight.shape[0]) + if convrot_groupsize in (16, 64, 256) and hasattr(_C, "quantize_int4_rowwise_convrot64"): + qact, x_scale = quantize_int4_rowwise_convrot64(x2d, convrot_groupsize) + else: + h = _build_hadamard(convrot_groupsize, device=x2d.device, dtype=x2d.dtype) + x_rot = _rotate_activation(x2d, h, convrot_groupsize).contiguous() + qact, x_scale = quantize_int4_rowwise(x_rot) + out = int4_linear( + qact, + qweight, + x_scale, + wscales, + bias=bias, + out_dtype=x.dtype, + ) + return out[: x2d.shape[0]].reshape(*orig_shape[:-1], qweight.shape[0]) + + def quantize_int8_rowwise_convrot( x_2d: torch.Tensor, group_size: int, @@ -1821,6 +2225,66 @@ def _build_constraints() -> dict: default_devices=cuda_devices, min_compute_capability=(8, 0), ), + "quantize_convrot_w4a4_weight": FunctionConstraints( + params={ + "weight": ParamConstraint( + dtypes=frozenset({torch.float32, torch.float16, torch.bfloat16}), + shape_rules=(ExactDims(2), DivisibleBy(dim=1, factor=64)), + ), + "convrot_groupsize": ParamConstraint(dtypes=frozenset({int})), + "quant_group_size": ParamConstraint(dtypes=frozenset({int})), + "stochastic_rounding": ParamConstraint(dtypes=frozenset({int})), + }, + default_devices=cuda_devices, + min_compute_capability=(8, 0), + ), + "dequantize_convrot_w4a4_weight": FunctionConstraints( + params={ + "qdata": ParamConstraint( + dtypes=frozenset({torch.int8}), + shape_rules=(ExactDims(2), DivisibleBy(dim=1, factor=32)), + ), + "scales": ParamConstraint( + dtypes=frozenset({torch.float32, torch.float16, torch.bfloat16}), + shape_rules=(ExactDims(1),), + ), + "convrot_groupsize": ParamConstraint(dtypes=frozenset({int})), + "quant_group_size": ParamConstraint(dtypes=frozenset({int})), + "output_dtype": ParamConstraint(dtypes=frozenset({torch.float32, torch.float16, torch.bfloat16})), + }, + default_devices=cuda_devices, + min_compute_capability=(8, 0), + ), + "convrot_w4a4_linear": FunctionConstraints( + params={ + "x": ParamConstraint( + dtypes=frozenset({torch.float32, torch.float16, torch.bfloat16}), + shape_rules=(MinDims(2),), + ), + "qweight": ParamConstraint( + dtypes=frozenset({torch.int8}), + shape_rules=(ExactDims(2), DivisibleBy(dim=1, factor=32)), + ), + "wscales": ParamConstraint( + dtypes=frozenset({torch.float32, torch.float16, torch.bfloat16}), + shape_rules=(ExactDims(1),), + ), + "bias": ParamConstraint(dtypes=frozenset({torch.float32, torch.float16, torch.bfloat16})), + "convrot_groupsize": ParamConstraint(dtypes=frozenset({int})), + "quant_group_size": ParamConstraint(dtypes=frozenset({int})), + }, + default_devices=cuda_devices, + min_compute_capability=(8, 0), + ), + "prepare_int4_weight_for_int8_linear": FunctionConstraints( + params={ + "weight": ParamConstraint( + dtypes=frozenset({torch.int8}), + shape_rules=(ExactDims(2), DivisibleBy(dim=1, factor=32)), + ), + }, + default_devices=cuda_devices, + ), "gemv_awq_w4a16": FunctionConstraints( params={ "x": ParamConstraint( diff --git a/comfy_kitchen/backends/cuda/dlpack_bindings.cpp b/comfy_kitchen/backends/cuda/dlpack_bindings.cpp index b788ce9..dc538a5 100644 --- a/comfy_kitchen/backends/cuda/dlpack_bindings.cpp +++ b/comfy_kitchen/backends/cuda/dlpack_bindings.cpp @@ -750,6 +750,74 @@ extern "C" { uint64_t seed, cudaStream_t stream); + void launch_quantize_int4_rowwise_kernel( + const void* input, + void* output, + void* scales, + int64_t M, + int64_t K, + int input_dtype_code, + bool stochastic, + uint64_t seed, + cudaStream_t stream); + + void launch_quantize_int4_rowwise_convrot64_kernel( + const void* input, + void* output, + void* scales, + int64_t M, + int64_t K, + int group_size, + int input_dtype_code, + bool stochastic, + uint64_t seed, + cudaStream_t stream); + + void launch_quantize_int4_rowwise_convrot64_to_int8_kernel( + const void* input, + void* output, + void* scales, + int64_t M, + int64_t K, + int group_size, + int input_dtype_code, + bool stochastic, + uint64_t seed, + cudaStream_t stream); + + void launch_dequantize_int4_convrot64_kernel( + const void* input, + const void* scales, + void* output, + int64_t M, + int64_t K, + int64_t scale_size, + int group_size, + int output_dtype_code, + cudaStream_t stream); + + void launch_int4_linear_kernel( + const void* act, + const void* weight, + const void* x_scales, + const void* weight_scales, + const void* bias, + void* output, + int64_t M, + int64_t N, + int64_t K, + bool has_bias, + int output_dtype_code, + int bias_dtype_code, + cudaStream_t stream); + + void launch_unpack_int4_to_int8_kernel( + const void* input, + void* output, + int64_t rows, + int64_t K_half, + cudaStream_t stream); + bool launch_cutlass_int8_dequant( const void* A, const void* B, @@ -763,6 +831,19 @@ extern "C" { int out_dtype_code, cudaStream_t stream); + bool launch_cutlass_int4_dequant( + const void* A, + const void* B, + const void* xs, + const void* ws, + const void* bias, + void* D, + int64_t M, + int64_t N, + int64_t K, + int out_dtype_code, + cudaStream_t stream); + void launch_quantize_int8_rowwise_convrot_kernel( const void* input, void* output, @@ -938,6 +1019,243 @@ void quantize_int8_rowwise( stream); } +void quantize_int4_rowwise( + nb::ndarray, nb::device::cuda> input, + nb::ndarray, nb::device::cuda> output, + nb::ndarray, nb::device::cuda> scales, + bool stochastic, + uint64_t seed, + uintptr_t stream_ptr) { + + const int64_t M = input.shape(0); + const int64_t K = input.shape(1); + if (K % 64 != 0) { + throw std::runtime_error("INT4 rowwise quantization requires K divisible by 64"); + } + if (output.shape(0) != M || output.shape(1) != K / 2) { + throw std::runtime_error("INT4 rowwise quantization output shape mismatch"); + } + if (scales.shape(0) != M || scales.shape(1) != 1) { + throw std::runtime_error("INT4 rowwise quantization scale shape mismatch"); + } + const int input_dtype_code = map_dtype_to_code(input.dtype()); + if (input_dtype_code < 0 || input_dtype_code > 2) { + throw std::runtime_error("Unsupported input dtype for INT4 rowwise quantization"); + } + + cudaStream_t stream = reinterpret_cast(stream_ptr); + launch_quantize_int4_rowwise_kernel( + input.data(), + output.data(), + scales.data(), + M, + K, + input_dtype_code, + stochastic, + seed, + stream); +} + +void quantize_int4_rowwise_convrot64( + nb::ndarray, nb::device::cuda> input, + nb::ndarray, nb::device::cuda> output, + nb::ndarray, nb::device::cuda> scales, + int group_size, + bool stochastic, + uint64_t seed, + uintptr_t stream_ptr) { + + const int64_t M = input.shape(0); + const int64_t K = input.shape(1); + if (group_size != 16 && group_size != 64 && group_size != 256) { + throw std::runtime_error("INT4 ConvRot quantization requires group_size 16, 64, or 256"); + } + if (K % group_size != 0) { + throw std::runtime_error("INT4 ConvRot quantization requires K divisible by group_size"); + } + if (output.shape(0) != M || output.shape(1) != K / 2) { + throw std::runtime_error("INT4 ConvRot quantization output shape mismatch"); + } + if (scales.shape(0) != M || scales.shape(1) != 1) { + throw std::runtime_error("INT4 ConvRot quantization scale shape mismatch"); + } + const int input_dtype_code = map_dtype_to_code(input.dtype()); + if (input_dtype_code < 0 || input_dtype_code > 2) { + throw std::runtime_error("Unsupported input dtype for INT4 ConvRot quantization"); + } + + cudaStream_t stream = reinterpret_cast(stream_ptr); + launch_quantize_int4_rowwise_convrot64_kernel( + input.data(), + output.data(), + scales.data(), + M, + K, + group_size, + input_dtype_code, + stochastic, + seed, + stream); +} + +void quantize_int4_rowwise_convrot64_to_int8( + nb::ndarray, nb::device::cuda> input, + nb::ndarray, nb::device::cuda> output, + nb::ndarray, nb::device::cuda> scales, + int group_size, + bool stochastic, + uint64_t seed, + uintptr_t stream_ptr) { + + const int64_t M = input.shape(0); + const int64_t K = input.shape(1); + if (output.shape(0) != M || output.shape(1) != K) { + throw std::runtime_error("INT4 ConvRot fallback activation output shape mismatch"); + } + if (scales.shape(0) != M || scales.shape(1) != 1) { + throw std::runtime_error("INT4 ConvRot fallback activation scales must have shape [M, 1]"); + } + if (group_size != 256 || K % group_size != 0) { + throw std::runtime_error("INT4 ConvRot fallback activation quantization requires group_size 256 and divisible K"); + } + const int input_dtype_code = map_dtype_to_code(input.dtype()); + if (input_dtype_code < 0 || input_dtype_code > 2) { + throw std::runtime_error("Unsupported input dtype for INT4 ConvRot fallback activation quantization"); + } + + cudaStream_t stream = reinterpret_cast(stream_ptr); + launch_quantize_int4_rowwise_convrot64_to_int8_kernel( + input.data(), + output.data(), + scales.data(), + M, + K, + static_cast(group_size), + input_dtype_code, + stochastic, + seed, + stream); +} + +void dequantize_int4_convrot64( + nb::ndarray, nb::device::cuda> input, + nb::ndarray, nb::device::cuda> scales, + nb::ndarray, nb::device::cuda> output, + int group_size, + uintptr_t stream_ptr) { + + const int64_t M = input.shape(0); + const int64_t K = input.shape(1) * 2; + if (group_size != 16 && group_size != 64 && group_size != 256) { + throw std::runtime_error("INT4 ConvRot dequantization requires group_size 16, 64, or 256"); + } + if (K % group_size != 0) { + throw std::runtime_error("INT4 ConvRot dequantization requires K divisible by group_size"); + } + if (output.shape(0) != M || output.shape(1) != K) { + throw std::runtime_error("INT4 ConvRot dequantization output shape mismatch"); + } + if (scales.size() != 1 && scales.size() != static_cast(M)) { + throw std::runtime_error("INT4 ConvRot dequantization scale must be scalar or per-row"); + } + const int output_dtype_code = map_dtype_to_code(output.dtype()); + if (output_dtype_code < 0 || output_dtype_code > 2) { + throw std::runtime_error("Unsupported output dtype for INT4 ConvRot dequantization"); + } + + cudaStream_t stream = reinterpret_cast(stream_ptr); + launch_dequantize_int4_convrot64_kernel( + input.data(), + scales.data(), + output.data(), + M, + K, + static_cast(scales.size()), + group_size, + output_dtype_code, + stream); +} + +void int4_linear( + nb::ndarray, nb::device::cuda> act, + nb::ndarray, nb::device::cuda> weight, + nb::ndarray x_scales, + nb::ndarray weight_scales, + nb::ndarray bias, + nb::ndarray, nb::device::cuda> output, + int output_dtype_code, + uintptr_t stream_ptr) { + + const int64_t M = act.shape(0); + const int64_t K_half = act.shape(1); + const int64_t N = weight.shape(0); + if (weight.shape(1) != K_half) { + throw std::runtime_error("INT4 linear K dimensions do not match"); + } + const int64_t K = K_half * 2; + if (K % 64 != 0) { + throw std::runtime_error("INT4 linear requires K divisible by 64"); + } + if (x_scales.size() != static_cast(M)) { + throw std::runtime_error("INT4 linear x_scales must have one value per row"); + } + if (weight_scales.size() != static_cast(N)) { + throw std::runtime_error("INT4 linear weight_scales must have one value per output channel"); + } + if (output.shape(0) != M || output.shape(1) != N) { + throw std::runtime_error("INT4 linear output shape mismatch"); + } + const int out_dtype = map_dtype_to_code(output.dtype()); + if (out_dtype != output_dtype_code) { + throw std::runtime_error("INT4 linear output dtype code mismatch"); + } + if (output_dtype_code < 0 || output_dtype_code > 2) { + throw std::runtime_error("Unsupported output dtype for INT4 linear"); + } + + const bool has_bias = bias.size() > 0; + int bias_dtype_code = output_dtype_code; + if (has_bias) { + if (bias.size() != static_cast(N)) { + throw std::runtime_error("INT4 linear bias shape mismatch"); + } + bias_dtype_code = map_dtype_to_code(bias.dtype()); + if (bias_dtype_code < 0 || bias_dtype_code > 2) { + throw std::runtime_error("Unsupported bias dtype for INT4 linear"); + } + } + + cudaStream_t stream = reinterpret_cast(stream_ptr); + launch_int4_linear_kernel( + act.data(), + weight.data(), + x_scales.data(), + weight_scales.data(), + has_bias ? bias.data() : nullptr, + output.data(), + M, + N, + K, + has_bias, + output_dtype_code, + bias_dtype_code, + stream); +} + +void unpack_int4_to_int8( + nb::ndarray, nb::device::cuda> input, + nb::ndarray, nb::device::cuda> output, + uintptr_t stream_ptr) { + + const int64_t rows = input.shape(0); + const int64_t K_half = input.shape(1); + if (output.shape(0) != rows || output.shape(1) != K_half * 2) { + throw std::runtime_error("unpack_int4_to_int8 output shape mismatch"); + } + cudaStream_t stream = reinterpret_cast(stream_ptr); + launch_unpack_int4_to_int8_kernel(input.data(), output.data(), rows, K_half, stream); +} + // INT8 GEMM + fused dequant (D = acc * xs[m] * ws[n] + bias[n]) via CUTLASS. // Returns true on success; false means caller falls back to cuBLAS + dequant. bool cutlass_int8_dequant( @@ -960,6 +1278,32 @@ bool cutlass_int8_dequant( bias_ptr, d.data(), M, N, K, out_dtype_code, stream); } +// INT4 GEMM + fused dequant via CUTLASS. A and B are packed signed int4 in int8 storage. +// Returns true on success; false means caller falls back to the hand-written int4 kernel. +bool cutlass_int4_dequant( + nb::ndarray, nb::device::cuda> a, // [M, K / 2] + nb::ndarray, nb::device::cuda> b, // [N, K / 2] + nb::ndarray xs, // [M] per-row act scale + nb::ndarray ws, // [N] per-col weight scale + nb::ndarray bias, // [N] float or empty + nb::ndarray, nb::device::cuda> d, // [M, N] output + int out_dtype_code, + uintptr_t stream_ptr) { + const int64_t M = a.shape(0); + const int64_t K_half = a.shape(1); + const int64_t N = b.shape(0); + if (b.shape(1) != K_half) throw std::runtime_error("cutlass_int4_dequant: K mismatch"); + if (d.shape(0) != M || d.shape(1) != N) throw std::runtime_error("cutlass_int4_dequant: D shape mismatch"); + if (xs.size() != static_cast(M)) throw std::runtime_error("cutlass_int4_dequant: xs shape mismatch"); + if (ws.size() != static_cast(N)) throw std::runtime_error("cutlass_int4_dequant: ws shape mismatch"); + const int64_t K = K_half * 2; + if (K % 64 != 0) throw std::runtime_error("cutlass_int4_dequant: K must be divisible by 64"); + cudaStream_t stream = reinterpret_cast(stream_ptr); + const void* bias_ptr = bias.size() > 0 ? bias.data() : nullptr; + return launch_cutlass_int4_dequant(a.data(), b.data(), xs.data(), ws.data(), + bias_ptr, d.data(), M, N, K, out_dtype_code, stream); +} + void quantize_int8_rowwise_convrot( nb::ndarray, nb::device::cuda> input, nb::ndarray, nb::device::cuda> output, @@ -1449,6 +1793,60 @@ NB_MODULE(_C, m) { nb::arg("seed"), nb::arg("stream_ptr")); + m.def("quantize_int4_rowwise", &quantize_int4_rowwise, + "Rowwise signed INT4 quantization for CUDA activations/weights", + nb::arg("input"), + nb::arg("output"), + nb::arg("scales"), + nb::arg("stochastic"), + nb::arg("seed"), + nb::arg("stream_ptr")); + + m.def("quantize_int4_rowwise_convrot64", &quantize_int4_rowwise_convrot64, + "Fused regular ConvRot-256 activation rotation plus rowwise signed INT4 quantization", + nb::arg("input"), + nb::arg("output"), + nb::arg("scales"), + nb::arg("group_size"), + nb::arg("stochastic"), + nb::arg("seed"), + nb::arg("stream_ptr")); + + m.def("quantize_int4_rowwise_convrot64_to_int8", &quantize_int4_rowwise_convrot64_to_int8, + "Fused ConvRot-256 activation rotation plus rowwise INT4-scale quantization into INT8 storage", + nb::arg("input"), + nb::arg("output"), + nb::arg("scales"), + nb::arg("group_size"), + nb::arg("stochastic"), + nb::arg("seed"), + nb::arg("stream_ptr")); + + m.def("dequantize_int4_convrot64", &dequantize_int4_convrot64, + "Fused packed signed INT4 dequantization plus regular ConvRot-256 inverse rotation", + nb::arg("input"), + nb::arg("scales"), + nb::arg("output"), + nb::arg("group_size"), + nb::arg("stream_ptr")); + + m.def("int4_linear", &int4_linear, + "Signed INT4 GEMM with rowwise x colwise dequantization, bias, and output cast", + nb::arg("act"), + nb::arg("weight"), + nb::arg("x_scales"), + nb::arg("weight_scales"), + nb::arg("bias"), + nb::arg("output"), + nb::arg("output_dtype_code"), + nb::arg("stream_ptr")); + + m.def("unpack_int4_to_int8", &unpack_int4_to_int8, + "Unpack row-major packed signed INT4 matrix to row-major INT8 matrix", + nb::arg("input"), + nb::arg("output"), + nb::arg("stream_ptr")); + m.def("cutlass_int8_dequant", &cutlass_int8_dequant, "INT8 GEMM + fused rowwise x colwise dequant + bias via CUTLASS; false -> fall back to cuBLAS", nb::arg("a"), @@ -1460,6 +1858,17 @@ NB_MODULE(_C, m) { nb::arg("out_dtype_code"), nb::arg("stream_ptr")); + m.def("cutlass_int4_dequant", &cutlass_int4_dequant, + "INT4 GEMM + fused rowwise x colwise dequant + bias via CUTLASS; false -> fall back to hand kernel", + nb::arg("a"), + nb::arg("b"), + nb::arg("xs"), + nb::arg("ws"), + nb::arg("bias"), + nb::arg("d"), + nb::arg("out_dtype_code"), + nb::arg("stream_ptr")); + m.def("quantize_int8_rowwise_convrot", &quantize_int8_rowwise_convrot, "Fused ConvRot Hadamard rotation + rowwise INT8 quantization", nb::arg("input"), diff --git a/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu b/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu new file mode 100644 index 0000000..94703e2 --- /dev/null +++ b/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu @@ -0,0 +1,2650 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Plain ConvRot W4A4 helpers: rowwise signed int4 quantization and int4 MMA +// with row/column dequant scales. +#include "svdquant_utils.cuh" + +#include +#include +#include +#include +#include + +#ifdef COMFY_HAVE_CUTLASS +#include "cutlass/cutlass.h" +#include "cutlass/gemm/device/gemm.h" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/default_gemm_universal_with_visitor.h" +#include "cutlass/epilogue/threadblock/fusion/visitors.hpp" +#include +#include +#include +#endif + +namespace { + +using comfy::svdquant::cp_async_16b; +using comfy::svdquant::cp_async_commit_group; +using comfy::svdquant::cp_async_wait_group; +using comfy::svdquant::kGroupSize; +using comfy::svdquant::kInt4Max; +using comfy::svdquant::mma_m16n8k64_s4s4s32; +using comfy::svdquant::pack_int4_pair; + +constexpr int kThreadsPerWarp = 32; +constexpr int kConvRotGroup = 256; + +__device__ __forceinline__ uint32_t pcg_hash(uint32_t x) { + const uint32_t state = x * 747796405u + 2891336453u; + const uint32_t word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u; + return (word >> 22u) ^ word; +} + +__device__ __forceinline__ float stochastic_rng_value(int64_t idx, uint64_t seed) { + const uint64_t key = static_cast(idx) + seed; + const uint32_t folded = static_cast(key) ^ static_cast(key >> 32); + return static_cast(pcg_hash(folded) >> 8) * 0x1.0p-24f; +} + +template +__device__ __forceinline__ int quantize_int4_value(float x, float inv_scale, uint64_t seed, int64_t idx) { + const float scaled = x * inv_scale; + int q; + if constexpr (STOCHASTIC) { + q = static_cast(floorf(scaled + stochastic_rng_value(idx, seed))); + } else { + q = __float2int_rn(scaled); + } + return max(-kInt4Max, min(kInt4Max, q)); +} + +__device__ __forceinline__ int unpack_int4_nibble(uint32_t v) { + return static_cast((v & 0x0fu) ^ 0x08u) - 8; +} + +#ifdef COMFY_HAVE_CUTLASS +template +struct FusedInt4Gemm { + using ElementA = cutlass::int4b_t; + using ElementB = cutlass::int4b_t; + using ElementC = ElementOutput; + using ElementAcc = int32_t; + using ElementCompute = float; + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + static constexpr int AlignA = 32; + static constexpr int AlignB = 32; + static constexpr int AlignC = 128 / cutlass::sizeof_bits::value; + using TB = cutlass::gemm::GemmShape; + using Warp = cutlass::gemm::GemmShape; + using Inst = cutlass::gemm::GemmShape<16, 8, 64>; + static constexpr int EVTStages = 1; + + using ThreadMap = cutlass::epilogue::threadblock::OutputTileThreadLayout; + using Accum = cutlass::epilogue::threadblock::VisitorAccFetch; + using XScale = cutlass::epilogue::threadblock::VisitorColBroadcast>; + using WScale = cutlass::epilogue::threadblock::VisitorRowBroadcast>; + using Bias = cutlass::epilogue::threadblock::VisitorRowBroadcast>; + using Mul0 = cutlass::epilogue::threadblock::VisitorCompute; + using EVT0 = cutlass::epilogue::threadblock::Sm80EVT; + using Mul1 = cutlass::epilogue::threadblock::VisitorCompute; + using EVT1 = cutlass::epilogue::threadblock::Sm80EVT; + using Add2 = cutlass::epilogue::threadblock::VisitorCompute; + using EVT2 = cutlass::epilogue::threadblock::Sm80EVT; + using StoreD = cutlass::epilogue::threadblock::VisitorAuxStore>; + using EVTD = cutlass::epilogue::threadblock::Sm80EVT; + + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor< + ElementA, LayoutA, cutlass::ComplexTransform::kNone, AlignA, + ElementB, LayoutB, cutlass::ComplexTransform::kNone, AlignB, + ElementC, LayoutC, AlignC, + ElementAcc, ElementCompute, + cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89, + TB, Warp, Inst, EVTD, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, + NumStages, cutlass::arch::OpMultiplyAddSaturate, EVTStages>::GemmKernel; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + static bool run(const int8_t* A, const int8_t* B, const float* xs, const float* ws, + const float* bias, ElementOutput* D, int M, int N, int K, cudaStream_t stream) { + cutlass::gemm::GemmCoord problem(M, N, K); + typename EVTD::Arguments cb{ + { { { {}, {const_cast(xs), 0.f, {cute::_1{}, cute::_0{}, M}}, {} }, + {const_cast(ws), 0.f, {cute::_0{}, cute::_1{}, N}}, {} }, + {const_cast(bias), 0.f, {cute::_0{}, cute::_1{}, N}}, {} }, + {D, {N, cute::_1{}, M * N}} }; + typename Gemm::Arguments args( + cutlass::gemm::GemmUniversalMode::kGemm, problem, 1, cb, + reinterpret_cast(const_cast(A)), + reinterpret_cast(const_cast(B)), + nullptr, nullptr, + (int64_t)M * K, (int64_t)N * K, 0, 0, K, K, 0, 0); + + Gemm gemm; + if (gemm.can_implement(args) != cutlass::Status::kSuccess) return false; + if (Gemm::get_workspace_size(args) != 0) return false; + if (gemm.initialize(args, nullptr, stream) != cutlass::Status::kSuccess) return false; + return gemm(stream) == cutlass::Status::kSuccess; + } +}; + +template +struct FusedInt4GemmNoBias { + using ElementA = cutlass::int4b_t; + using ElementB = cutlass::int4b_t; + using ElementC = ElementOutput; + using ElementAcc = int32_t; + using ElementCompute = float; + using LayoutA = cutlass::layout::RowMajor; + using LayoutB = cutlass::layout::ColumnMajor; + using LayoutC = cutlass::layout::RowMajor; + static constexpr int AlignA = 32; + static constexpr int AlignB = 32; + static constexpr int AlignC = 128 / cutlass::sizeof_bits::value; + using TB = cutlass::gemm::GemmShape; + using Warp = cutlass::gemm::GemmShape; + using Inst = cutlass::gemm::GemmShape<16, 8, 64>; + static constexpr int EVTStages = 1; + + using ThreadMap = cutlass::epilogue::threadblock::OutputTileThreadLayout; + using Accum = cutlass::epilogue::threadblock::VisitorAccFetch; + using XScale = cutlass::epilogue::threadblock::VisitorColBroadcast>; + using WScale = cutlass::epilogue::threadblock::VisitorRowBroadcast>; + using Mul0 = cutlass::epilogue::threadblock::VisitorCompute; + using EVT0 = cutlass::epilogue::threadblock::Sm80EVT; + using Mul1 = cutlass::epilogue::threadblock::VisitorCompute; + using EVT1 = cutlass::epilogue::threadblock::Sm80EVT; + using StoreD = cutlass::epilogue::threadblock::VisitorAuxStore>; + using EVTD = cutlass::epilogue::threadblock::Sm80EVT; + + using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor< + ElementA, LayoutA, cutlass::ComplexTransform::kNone, AlignA, + ElementB, LayoutB, cutlass::ComplexTransform::kNone, AlignB, + ElementC, LayoutC, AlignC, + ElementAcc, ElementCompute, + cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89, + TB, Warp, Inst, EVTD, + cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, + NumStages, cutlass::arch::OpMultiplyAddSaturate, EVTStages>::GemmKernel; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + + static bool run(const int8_t* A, const int8_t* B, const float* xs, const float* ws, + ElementOutput* D, int M, int N, int K, cudaStream_t stream) { + cutlass::gemm::GemmCoord problem(M, N, K); + typename EVTD::Arguments cb{ + { { {}, {const_cast(xs), 0.f, {cute::_1{}, cute::_0{}, M}}, {} }, + {const_cast(ws), 0.f, {cute::_0{}, cute::_1{}, N}}, {} }, + {D, {N, cute::_1{}, M * N}} }; + typename Gemm::Arguments args( + cutlass::gemm::GemmUniversalMode::kGemm, problem, 1, cb, + reinterpret_cast(const_cast(A)), + reinterpret_cast(const_cast(B)), + nullptr, nullptr, + (int64_t)M * K, (int64_t)N * K, 0, 0, K, K, 0, 0); + + Gemm gemm; + if (gemm.can_implement(args) != cutlass::Status::kSuccess) return false; + if (Gemm::get_workspace_size(args) != 0) return false; + if (gemm.initialize(args, nullptr, stream) != cutlass::Status::kSuccess) return false; + return gemm(stream) == cutlass::Status::kSuccess; + } +}; + +template +bool dispatch_fused_int4(const int8_t* A, const int8_t* B, const float* xs, const float* ws, + const float* bias, OutT* D, int M, int N, int K, cudaStream_t stream) { + using Fn = bool (*)(const int8_t*, const int8_t*, const float*, const float*, const float*, OutT*, int, int, int, cudaStream_t); + static const Fn runners[] = { + &FusedInt4Gemm::run, + &FusedInt4Gemm::run, + &FusedInt4Gemm::run, + &FusedInt4Gemm::run, + &FusedInt4Gemm::run, + &FusedInt4Gemm::run, + &FusedInt4Gemm::run, + }; + constexpr int NC = sizeof(runners) / sizeof(runners[0]); + + if (M == 4608 && N == 3072 && K == 15360) { + return runners[0](A, B, xs, ws, bias, D, M, N, K, stream); + } + + static std::mutex mtx; + static std::map, int> cache; + const std::tuple key{M, N, K}; + + static thread_local int last_m = -1; + static thread_local int last_n = -1; + static thread_local int last_k = -1; + static thread_local int last_best = -2; + if (M == last_m && N == last_n && K == last_k) { + if (last_best < 0) return false; + if (runners[last_best](A, B, xs, ws, bias, D, M, N, K, stream)) return true; + return runners[last_best](A, B, xs, ws, bias, D, M, N, K, stream); + } + + int best; + { + std::lock_guard lk(mtx); + auto it = cache.find(key); + best = (it != cache.end()) ? it->second : -2; + } + if (best == -2) { + best = -1; + float best_ms = 1e30f; + cudaEvent_t s, e; + cudaEventCreate(&s); + cudaEventCreate(&e); + for (int i = 0; i < NC; ++i) { + if (!runners[i](A, B, xs, ws, bias, D, M, N, K, stream)) continue; + cudaStreamSynchronize(stream); + for (int r = 0; r < 8; ++r) { + runners[i](A, B, xs, ws, bias, D, M, N, K, stream); + } + cudaStreamSynchronize(stream); + cudaEventRecord(s, stream); + for (int r = 0; r < 32; ++r) { + runners[i](A, B, xs, ws, bias, D, M, N, K, stream); + } + cudaEventRecord(e, stream); + cudaEventSynchronize(e); + float ms = 0.f; + cudaEventElapsedTime(&ms, s, e); + if (ms < best_ms) { + best_ms = ms; + best = i; + } + } + cudaEventDestroy(s); + cudaEventDestroy(e); + std::lock_guard lk(mtx); + cache[key] = best; + } + last_m = M; + last_n = N; + last_k = K; + last_best = best; + if (best < 0) return false; + if (runners[best](A, B, xs, ws, bias, D, M, N, K, stream)) return true; + return runners[best](A, B, xs, ws, bias, D, M, N, K, stream); +} + +template +bool dispatch_fused_int4_no_bias(const int8_t* A, const int8_t* B, const float* xs, const float* ws, + OutT* D, int M, int N, int K, cudaStream_t stream) { + using Fn = bool (*)(const int8_t*, const int8_t*, const float*, const float*, OutT*, int, int, int, cudaStream_t); + static const Fn runners[] = { + &FusedInt4GemmNoBias::run, + &FusedInt4GemmNoBias::run, + &FusedInt4GemmNoBias::run, + &FusedInt4GemmNoBias::run, + &FusedInt4GemmNoBias::run, + &FusedInt4GemmNoBias::run, + &FusedInt4GemmNoBias::run, + }; + constexpr int NC = sizeof(runners) / sizeof(runners[0]); + + if (M == 4608 && N == 3072 && K == 15360) { + return runners[0](A, B, xs, ws, D, M, N, K, stream); + } + + static std::mutex mtx; + static std::map, int> cache; + const std::tuple key{M, N, K}; + + static thread_local int last_m = -1; + static thread_local int last_n = -1; + static thread_local int last_k = -1; + static thread_local int last_best = -2; + if (M == last_m && N == last_n && K == last_k) { + if (last_best < 0) return false; + if (runners[last_best](A, B, xs, ws, D, M, N, K, stream)) return true; + return runners[last_best](A, B, xs, ws, D, M, N, K, stream); + } + + int best; + { + std::lock_guard lk(mtx); + auto it = cache.find(key); + best = (it != cache.end()) ? it->second : -2; + } + if (best == -2) { + best = -1; + float best_ms = 1e30f; + cudaEvent_t s, e; + cudaEventCreate(&s); + cudaEventCreate(&e); + for (int i = 0; i < NC; ++i) { + if (!runners[i](A, B, xs, ws, D, M, N, K, stream)) continue; + cudaStreamSynchronize(stream); + for (int r = 0; r < 8; ++r) { + runners[i](A, B, xs, ws, D, M, N, K, stream); + } + cudaStreamSynchronize(stream); + cudaEventRecord(s, stream); + for (int r = 0; r < 32; ++r) { + runners[i](A, B, xs, ws, D, M, N, K, stream); + } + cudaEventRecord(e, stream); + cudaEventSynchronize(e); + float ms = 0.f; + cudaEventElapsedTime(&ms, s, e); + if (ms < best_ms) { + best_ms = ms; + best = i; + } + } + cudaEventDestroy(s); + cudaEventDestroy(e); + std::lock_guard lk(mtx); + cache[key] = best; + } + last_m = M; + last_n = N; + last_k = K; + last_best = best; + if (best < 0) return false; + if (runners[best](A, B, xs, ws, D, M, N, K, stream)) return true; + return runners[best](A, B, xs, ws, D, M, N, K, stream); +} + +#endif + +template +__device__ __forceinline__ float to_float(T v); + +template<> +__device__ __forceinline__ float to_float(float v) { return v; } + +template<> +__device__ __forceinline__ float to_float<__half>(__half v) { return __half2float(v); } + +template<> +__device__ __forceinline__ float to_float<__nv_bfloat16>(__nv_bfloat16 v) { return __bfloat162float(v); } + +template +__device__ __forceinline__ T from_float(float v); + +template<> +__device__ __forceinline__ float from_float(float v) { return v; } + +template<> +__device__ __forceinline__ __half from_float<__half>(float v) { return __float2half(v); } + +template<> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { return __float2bfloat16(v); } + +template +__device__ __forceinline__ float block_reduce_max(float v, float* warp_smem, float* block_smem) { + const int lane = threadIdx.x & (kThreadsPerWarp - 1); + const int wid = threadIdx.x >> 5; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + v = fmaxf(v, __shfl_down_sync(0xffffffffu, v, offset)); + } + if (lane == 0) { + warp_smem[wid] = v; + } + __syncthreads(); + if (wid == 0) { + float total = lane < NUM_WARPS ? warp_smem[lane] : 0.0f; + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + total = fmaxf(total, __shfl_down_sync(0xffffffffu, total, offset)); + } + if (lane == 0) { + *block_smem = total; + } + } + __syncthreads(); + return *block_smem; +} + +template +__device__ __forceinline__ void convrot_fht_stage64_typed( + const T* __restrict__ src, + T* __restrict__ dst, + int lane, + int stride) +{ + const int base = (lane % stride) + (lane / stride) * (4 * stride); + const float x0 = to_float(src[base]); + const float x1 = to_float(src[base + stride]); + const float x2 = to_float(src[base + 2 * stride]); + const float x3 = to_float(src[base + 3 * stride]); + dst[base] = from_float(0.5f * (x0 + x1 + x2 - x3)); + dst[base + stride] = from_float(0.5f * (x0 + x1 - x2 + x3)); + dst[base + 2 * stride] = from_float(0.5f * (x0 - x1 + x2 + x3)); + dst[base + 3 * stride] = from_float(0.5f * (-x0 + x1 + x2 + x3)); +} + +template +__device__ __forceinline__ void convrot_fht_stage64_store_typed( + const T* __restrict__ src, + T* __restrict__ output, + int lane, + int stride) +{ + const int base = (lane % stride) + (lane / stride) * (4 * stride); + const float x0 = to_float(src[base]); + const float x1 = to_float(src[base + stride]); + const float x2 = to_float(src[base + 2 * stride]); + const float x3 = to_float(src[base + 3 * stride]); + output[base] = from_float(0.5f * (x0 + x1 + x2 - x3)); + output[base + stride] = from_float(0.5f * (x0 + x1 - x2 + x3)); + output[base + 2 * stride] = from_float(0.5f * (x0 - x1 + x2 + x3)); + output[base + 3 * stride] = from_float(0.5f * (-x0 + x1 + x2 + x3)); +} + +template +__device__ __forceinline__ float convrot_fht_stage64_store_absmax_typed( + const T* __restrict__ src, + T* __restrict__ output, + int lane, + int stride) +{ + const int base = (lane % stride) + (lane / stride) * (4 * stride); + const float x0 = to_float(src[base]); + const float x1 = to_float(src[base + stride]); + const float x2 = to_float(src[base + 2 * stride]); + const float x3 = to_float(src[base + 3 * stride]); + const T y0_t = from_float(0.5f * (x0 + x1 + x2 - x3)); + const T y1_t = from_float(0.5f * (x0 + x1 - x2 + x3)); + const T y2_t = from_float(0.5f * (x0 - x1 + x2 + x3)); + const T y3_t = from_float(0.5f * (-x0 + x1 + x2 + x3)); + output[base] = y0_t; + output[base + stride] = y1_t; + output[base + 2 * stride] = y2_t; + output[base + 3 * stride] = y3_t; + const float y0 = to_float(y0_t); + const float y1 = to_float(y1_t); + const float y2 = to_float(y2_t); + const float y3 = to_float(y3_t); + return fmaxf(fmaxf(fabsf(y0), fabsf(y1)), fmaxf(fabsf(y2), fabsf(y3))); +} + +template +__device__ __forceinline__ void convrot_fht_stage64_float( + const float* __restrict__ src, + float* __restrict__ dst, + int lane) +{ + const int base = (lane % S) + (lane / S) * (4 * S); + const float x0 = src[base]; + const float x1 = src[base + S]; + const float x2 = src[base + 2 * S]; + const float x3 = src[base + 3 * S]; + dst[base] = 0.5f * (x0 + x1 + x2 - x3); + dst[base + S] = 0.5f * (x0 + x1 - x2 + x3); + dst[base + 2 * S] = 0.5f * (x0 - x1 + x2 + x3); + dst[base + 3 * S] = 0.5f * (-x0 + x1 + x2 + x3); +} + +template +__device__ __forceinline__ float convrot_fht_stage64_store_absmax_float( + const float* __restrict__ src, + float* __restrict__ output, + int lane) +{ + const int base = (lane % S) + (lane / S) * (4 * S); + const float x0 = src[base]; + const float x1 = src[base + S]; + const float x2 = src[base + 2 * S]; + const float x3 = src[base + 3 * S]; + const float y0 = 0.5f * (x0 + x1 + x2 - x3); + const float y1 = 0.5f * (x0 + x1 - x2 + x3); + const float y2 = 0.5f * (x0 - x1 + x2 + x3); + const float y3 = 0.5f * (-x0 + x1 + x2 + x3); + output[base] = y0; + output[base + S] = y1; + output[base + 2 * S] = y2; + output[base + 3 * S] = y3; + return fmaxf(fmaxf(fabsf(y0), fabsf(y1)), fmaxf(fabsf(y2), fabsf(y3))); +} + +__device__ __forceinline__ __half half_fht_scale(__half v) { + return __hmul(v, __float2half(0.5f)); +} + +__device__ __forceinline__ __nv_bfloat16 bfloat16_fht_scale(__nv_bfloat16 v) { + return __hmul(v, __float2bfloat16(0.5f)); +} + +template +__device__ __forceinline__ void convrot_h4_store_typed( + T x0, + T x1, + T x2, + T x3, + T* __restrict__ output, + int base) +{ + const float f0 = to_float(x0); + const float f1 = to_float(x1); + const float f2 = to_float(x2); + const float f3 = to_float(x3); + output[base] = from_float(0.5f * (f0 + f1 + f2 - f3)); + output[base + 1] = from_float(0.5f * (f0 + f1 - f2 + f3)); + output[base + 2] = from_float(0.5f * (f0 - f1 + f2 + f3)); + output[base + 3] = from_float(0.5f * (-f0 + f1 + f2 + f3)); +} + +__device__ __forceinline__ void convrot_h4_store_typed( + __half x0, + __half x1, + __half x2, + __half x3, + __half* __restrict__ output, + int base) +{ + const __half x01 = __hadd(x0, x1); + const __half x0m1 = __hsub(x0, x1); + const __half x23 = __hadd(x2, x3); + const __half x2m3 = __hsub(x2, x3); + output[base] = half_fht_scale(__hadd(x01, x2m3)); + output[base + 1] = half_fht_scale(__hsub(x01, x2m3)); + output[base + 2] = half_fht_scale(__hadd(x0m1, x23)); + output[base + 3] = half_fht_scale(__hsub(x23, x0m1)); +} + +__device__ __forceinline__ void convrot_h4_store_typed( + __nv_bfloat16 x0, + __nv_bfloat16 x1, + __nv_bfloat16 x2, + __nv_bfloat16 x3, + __nv_bfloat16* __restrict__ output, + int base) +{ + const __nv_bfloat16 x01 = __hadd(x0, x1); + const __nv_bfloat16 x0m1 = __hsub(x0, x1); + const __nv_bfloat16 x23 = __hadd(x2, x3); + const __nv_bfloat16 x2m3 = __hsub(x2, x3); + output[base] = bfloat16_fht_scale(__hadd(x01, x2m3)); + output[base + 1] = bfloat16_fht_scale(__hsub(x01, x2m3)); + output[base + 2] = bfloat16_fht_scale(__hadd(x0m1, x23)); + output[base + 3] = bfloat16_fht_scale(__hsub(x23, x0m1)); +} + +template +__device__ __forceinline__ void dequant_int4_h4_store_typed( + int q0, + int q1, + int q2, + int q3, + float scale, + T* __restrict__ output, + int base) +{ + const float first_stage_scale = 0.5f * scale; + output[base] = from_float(static_cast(q0 + q1 + q2 - q3) * first_stage_scale); + output[base + 1] = from_float(static_cast(q0 + q1 - q2 + q3) * first_stage_scale); + output[base + 2] = from_float(static_cast(q0 - q1 + q2 + q3) * first_stage_scale); + output[base + 3] = from_float(static_cast(-q0 + q1 + q2 + q3) * first_stage_scale); +} + +__device__ __forceinline__ void dequant_int4_h4_store_typed( + int q0, + int q1, + int q2, + int q3, + float scale, + __half* __restrict__ output, + int base) +{ + const __half2 s = __float2half2_rn(0.5f * scale); + const __half2 y01 = __hmul2( + __floats2half2_rn( + static_cast(q0 + q1 + q2 - q3), + static_cast(q0 + q1 - q2 + q3)), + s); + const __half2 y23 = __hmul2( + __floats2half2_rn( + static_cast(q0 - q1 + q2 + q3), + static_cast(-q0 + q1 + q2 + q3)), + s); + *reinterpret_cast<__half2*>(output + base) = y01; + *reinterpret_cast<__half2*>(output + base + 2) = y23; +} + +__device__ __forceinline__ void dequant_int4_h4_store_typed( + int q0, + int q1, + int q2, + int q3, + float scale, + __nv_bfloat16* __restrict__ output, + int base) +{ + const __nv_bfloat162 s = __float2bfloat162_rn(0.5f * scale); + const __nv_bfloat162 y01 = __hmul2( + __floats2bfloat162_rn( + static_cast(q0 + q1 + q2 - q3), + static_cast(q0 + q1 - q2 + q3)), + s); + const __nv_bfloat162 y23 = __hmul2( + __floats2bfloat162_rn( + static_cast(q0 - q1 + q2 + q3), + static_cast(-q0 + q1 + q2 + q3)), + s); + *reinterpret_cast<__nv_bfloat162*>(output + base) = y01; + *reinterpret_cast<__nv_bfloat162*>(output + base + 2) = y23; +} + +__device__ __forceinline__ void convrot_fht_stage64_typed( + const __half* __restrict__ src, + __half* __restrict__ dst, + int lane, + int stride) +{ + const int base = (lane % stride) + (lane / stride) * (4 * stride); + const __half x0 = src[base]; + const __half x1 = src[base + stride]; + const __half x2 = src[base + 2 * stride]; + const __half x3 = src[base + 3 * stride]; + const __half x01 = __hadd(x0, x1); + const __half x0m1 = __hsub(x0, x1); + const __half x23 = __hadd(x2, x3); + const __half x2m3 = __hsub(x2, x3); + dst[base] = half_fht_scale(__hadd(x01, x2m3)); + dst[base + stride] = half_fht_scale(__hsub(x01, x2m3)); + dst[base + 2 * stride] = half_fht_scale(__hadd(x0m1, x23)); + dst[base + 3 * stride] = half_fht_scale(__hsub(x23, x0m1)); +} + +__device__ __forceinline__ void convrot_fht_stage64_typed( + const __nv_bfloat16* __restrict__ src, + __nv_bfloat16* __restrict__ dst, + int lane, + int stride) +{ + const int base = (lane % stride) + (lane / stride) * (4 * stride); + const __nv_bfloat16 x0 = src[base]; + const __nv_bfloat16 x1 = src[base + stride]; + const __nv_bfloat16 x2 = src[base + 2 * stride]; + const __nv_bfloat16 x3 = src[base + 3 * stride]; + const __nv_bfloat16 x01 = __hadd(x0, x1); + const __nv_bfloat16 x0m1 = __hsub(x0, x1); + const __nv_bfloat16 x23 = __hadd(x2, x3); + const __nv_bfloat16 x2m3 = __hsub(x2, x3); + dst[base] = bfloat16_fht_scale(__hadd(x01, x2m3)); + dst[base + stride] = bfloat16_fht_scale(__hsub(x01, x2m3)); + dst[base + 2 * stride] = bfloat16_fht_scale(__hadd(x0m1, x23)); + dst[base + 3 * stride] = bfloat16_fht_scale(__hsub(x23, x0m1)); +} + +template +__device__ __forceinline__ void convrot_fht_stage64_vec2_typed( + const T* __restrict__ src, + T* __restrict__ dst, + int lane, + int stride) +{ + convrot_fht_stage64_typed(src, dst, lane, stride); +} + +__device__ __forceinline__ void convrot_fht_stage64_vec2_typed( + const __half* __restrict__ src, + __half* __restrict__ dst, + int lane, + int stride) +{ + if (lane & 1) { + return; + } + const int base = (lane % stride) + (lane / stride) * (4 * stride); + const __half2 x0 = *reinterpret_cast(src + base); + const __half2 x1 = *reinterpret_cast(src + base + stride); + const __half2 x2 = *reinterpret_cast(src + base + 2 * stride); + const __half2 x3 = *reinterpret_cast(src + base + 3 * stride); + const __half2 x01 = __hadd2(x0, x1); + const __half2 x0m1 = __hsub2(x0, x1); + const __half2 x23 = __hadd2(x2, x3); + const __half2 x2m3 = __hsub2(x2, x3); + const __half2 s = __float2half2_rn(0.5f); + *reinterpret_cast<__half2*>(dst + base) = __hmul2(__hadd2(x01, x2m3), s); + *reinterpret_cast<__half2*>(dst + base + stride) = __hmul2(__hsub2(x01, x2m3), s); + *reinterpret_cast<__half2*>(dst + base + 2 * stride) = __hmul2(__hadd2(x0m1, x23), s); + *reinterpret_cast<__half2*>(dst + base + 3 * stride) = __hmul2(__hsub2(x23, x0m1), s); +} + +__device__ __forceinline__ void convrot_fht_stage64_vec2_typed( + const __nv_bfloat16* __restrict__ src, + __nv_bfloat16* __restrict__ dst, + int lane, + int stride) +{ + if (lane & 1) { + return; + } + const int base = (lane % stride) + (lane / stride) * (4 * stride); + const __nv_bfloat162 x0 = *reinterpret_cast(src + base); + const __nv_bfloat162 x1 = *reinterpret_cast(src + base + stride); + const __nv_bfloat162 x2 = *reinterpret_cast(src + base + 2 * stride); + const __nv_bfloat162 x3 = *reinterpret_cast(src + base + 3 * stride); + const __nv_bfloat162 x01 = __hadd2(x0, x1); + const __nv_bfloat162 x0m1 = __hsub2(x0, x1); + const __nv_bfloat162 x23 = __hadd2(x2, x3); + const __nv_bfloat162 x2m3 = __hsub2(x2, x3); + const __nv_bfloat162 s = __float2bfloat162_rn(0.5f); + *reinterpret_cast<__nv_bfloat162*>(dst + base) = __hmul2(__hadd2(x01, x2m3), s); + *reinterpret_cast<__nv_bfloat162*>(dst + base + stride) = __hmul2(__hsub2(x01, x2m3), s); + *reinterpret_cast<__nv_bfloat162*>(dst + base + 2 * stride) = __hmul2(__hadd2(x0m1, x23), s); + *reinterpret_cast<__nv_bfloat162*>(dst + base + 3 * stride) = __hmul2(__hsub2(x23, x0m1), s); +} + +__device__ __forceinline__ void convrot_fht_stage64_store_typed( + const __half* __restrict__ src, + __half* __restrict__ output, + int lane, + int stride) +{ + convrot_fht_stage64_typed(src, output, lane, stride); +} + +__device__ __forceinline__ void convrot_fht_stage64_store_typed( + const __nv_bfloat16* __restrict__ src, + __nv_bfloat16* __restrict__ output, + int lane, + int stride) +{ + convrot_fht_stage64_typed(src, output, lane, stride); +} + +template +__device__ __forceinline__ void convrot_fht_stage64_store_final_typed( + const T* __restrict__ src, + T* __restrict__ output, + int lane) +{ + convrot_fht_stage64_store_typed(src, output, lane, 64); +} + +__device__ __forceinline__ void convrot_fht_stage64_store_final_typed( + const __half* __restrict__ src, + __half* __restrict__ output, + int lane) +{ + if (lane >= 32) { + return; + } + const int col = lane * 2; + const __half2 x0 = *reinterpret_cast(src + col); + const __half2 x1 = *reinterpret_cast(src + col + 64); + const __half2 x2 = *reinterpret_cast(src + col + 128); + const __half2 x3 = *reinterpret_cast(src + col + 192); + const __half2 x01 = __hadd2(x0, x1); + const __half2 x0m1 = __hsub2(x0, x1); + const __half2 x23 = __hadd2(x2, x3); + const __half2 x2m3 = __hsub2(x2, x3); + const __half2 s = __float2half2_rn(0.5f); + *reinterpret_cast<__half2*>(output + col) = __hmul2(__hadd2(x01, x2m3), s); + *reinterpret_cast<__half2*>(output + col + 64) = __hmul2(__hsub2(x01, x2m3), s); + *reinterpret_cast<__half2*>(output + col + 128) = __hmul2(__hadd2(x0m1, x23), s); + *reinterpret_cast<__half2*>(output + col + 192) = __hmul2(__hsub2(x23, x0m1), s); +} + +__device__ __forceinline__ void convrot_fht_stage64_store_final_typed( + const __nv_bfloat16* __restrict__ src, + __nv_bfloat16* __restrict__ output, + int lane) +{ + if (lane >= 32) { + return; + } + const int col = lane * 2; + const __nv_bfloat162 x0 = *reinterpret_cast(src + col); + const __nv_bfloat162 x1 = *reinterpret_cast(src + col + 64); + const __nv_bfloat162 x2 = *reinterpret_cast(src + col + 128); + const __nv_bfloat162 x3 = *reinterpret_cast(src + col + 192); + const __nv_bfloat162 x01 = __hadd2(x0, x1); + const __nv_bfloat162 x0m1 = __hsub2(x0, x1); + const __nv_bfloat162 x23 = __hadd2(x2, x3); + const __nv_bfloat162 x2m3 = __hsub2(x2, x3); + const __nv_bfloat162 s = __float2bfloat162_rn(0.5f); + *reinterpret_cast<__nv_bfloat162*>(output + col) = __hmul2(__hadd2(x01, x2m3), s); + *reinterpret_cast<__nv_bfloat162*>(output + col + 64) = __hmul2(__hsub2(x01, x2m3), s); + *reinterpret_cast<__nv_bfloat162*>(output + col + 128) = __hmul2(__hadd2(x0m1, x23), s); + *reinterpret_cast<__nv_bfloat162*>(output + col + 192) = __hmul2(__hsub2(x23, x0m1), s); +} + +__device__ __forceinline__ float convrot_fht_stage64_store_absmax_typed( + const __half* __restrict__ src, + __half* __restrict__ output, + int lane, + int stride) +{ + if (stride == 64) { + if (lane >= 32) { + return 0.0f; + } + const int col = lane * 2; + const __half2 x0 = *reinterpret_cast(src + col); + const __half2 x1 = *reinterpret_cast(src + col + 64); + const __half2 x2 = *reinterpret_cast(src + col + 128); + const __half2 x3 = *reinterpret_cast(src + col + 192); + const __half2 x01 = __hadd2(x0, x1); + const __half2 x0m1 = __hsub2(x0, x1); + const __half2 x23 = __hadd2(x2, x3); + const __half2 x2m3 = __hsub2(x2, x3); + const __half2 s = __float2half2_rn(0.5f); + const __half2 y0 = __hmul2(__hadd2(x01, x2m3), s); + const __half2 y1 = __hmul2(__hsub2(x01, x2m3), s); + const __half2 y2 = __hmul2(__hadd2(x0m1, x23), s); + const __half2 y3 = __hmul2(__hsub2(x23, x0m1), s); + *reinterpret_cast<__half2*>(output + col) = y0; + *reinterpret_cast<__half2*>(output + col + 64) = y1; + *reinterpret_cast<__half2*>(output + col + 128) = y2; + *reinterpret_cast<__half2*>(output + col + 192) = y3; + const float2 f0 = __half22float2(y0); + const float2 f1 = __half22float2(y1); + const float2 f2 = __half22float2(y2); + const float2 f3 = __half22float2(y3); + float v = fmaxf(fabsf(f0.x), fabsf(f0.y)); + v = fmaxf(v, fmaxf(fabsf(f1.x), fabsf(f1.y))); + v = fmaxf(v, fmaxf(fabsf(f2.x), fabsf(f2.y))); + v = fmaxf(v, fmaxf(fabsf(f3.x), fabsf(f3.y))); + return v; + } + const int base = (lane % stride) + (lane / stride) * (4 * stride); + const __half x0 = src[base]; + const __half x1 = src[base + stride]; + const __half x2 = src[base + 2 * stride]; + const __half x3 = src[base + 3 * stride]; + const __half x01 = __hadd(x0, x1); + const __half x0m1 = __hsub(x0, x1); + const __half x23 = __hadd(x2, x3); + const __half x2m3 = __hsub(x2, x3); + const __half y0 = half_fht_scale(__hadd(x01, x2m3)); + const __half y1 = half_fht_scale(__hsub(x01, x2m3)); + const __half y2 = half_fht_scale(__hadd(x0m1, x23)); + const __half y3 = half_fht_scale(__hsub(x23, x0m1)); + output[base] = y0; + output[base + stride] = y1; + output[base + 2 * stride] = y2; + output[base + 3 * stride] = y3; + const float f0 = fabsf(__half2float(y0)); + const float f1 = fabsf(__half2float(y1)); + const float f2 = fabsf(__half2float(y2)); + const float f3 = fabsf(__half2float(y3)); + return fmaxf(fmaxf(f0, f1), fmaxf(f2, f3)); +} + +__device__ __forceinline__ float convrot_fht_stage64_store_absmax_typed( + const __nv_bfloat16* __restrict__ src, + __nv_bfloat16* __restrict__ output, + int lane, + int stride) +{ + if (stride == 64) { + if (lane >= 32) { + return 0.0f; + } + const int col = lane * 2; + const __nv_bfloat162 x0 = *reinterpret_cast(src + col); + const __nv_bfloat162 x1 = *reinterpret_cast(src + col + 64); + const __nv_bfloat162 x2 = *reinterpret_cast(src + col + 128); + const __nv_bfloat162 x3 = *reinterpret_cast(src + col + 192); + const __nv_bfloat162 x01 = __hadd2(x0, x1); + const __nv_bfloat162 x0m1 = __hsub2(x0, x1); + const __nv_bfloat162 x23 = __hadd2(x2, x3); + const __nv_bfloat162 x2m3 = __hsub2(x2, x3); + const __nv_bfloat162 s = __float2bfloat162_rn(0.5f); + const __nv_bfloat162 y0 = __hmul2(__hadd2(x01, x2m3), s); + const __nv_bfloat162 y1 = __hmul2(__hsub2(x01, x2m3), s); + const __nv_bfloat162 y2 = __hmul2(__hadd2(x0m1, x23), s); + const __nv_bfloat162 y3 = __hmul2(__hsub2(x23, x0m1), s); + *reinterpret_cast<__nv_bfloat162*>(output + col) = y0; + *reinterpret_cast<__nv_bfloat162*>(output + col + 64) = y1; + *reinterpret_cast<__nv_bfloat162*>(output + col + 128) = y2; + *reinterpret_cast<__nv_bfloat162*>(output + col + 192) = y3; + const float2 f0 = __bfloat1622float2(y0); + const float2 f1 = __bfloat1622float2(y1); + const float2 f2 = __bfloat1622float2(y2); + const float2 f3 = __bfloat1622float2(y3); + float v = fmaxf(fabsf(f0.x), fabsf(f0.y)); + v = fmaxf(v, fmaxf(fabsf(f1.x), fabsf(f1.y))); + v = fmaxf(v, fmaxf(fabsf(f2.x), fabsf(f2.y))); + v = fmaxf(v, fmaxf(fabsf(f3.x), fabsf(f3.y))); + return v; + } + const int base = (lane % stride) + (lane / stride) * (4 * stride); + const __nv_bfloat16 x0 = src[base]; + const __nv_bfloat16 x1 = src[base + stride]; + const __nv_bfloat16 x2 = src[base + 2 * stride]; + const __nv_bfloat16 x3 = src[base + 3 * stride]; + const __nv_bfloat16 x01 = __hadd(x0, x1); + const __nv_bfloat16 x0m1 = __hsub(x0, x1); + const __nv_bfloat16 x23 = __hadd(x2, x3); + const __nv_bfloat16 x2m3 = __hsub(x2, x3); + const __nv_bfloat16 y0 = bfloat16_fht_scale(__hadd(x01, x2m3)); + const __nv_bfloat16 y1 = bfloat16_fht_scale(__hsub(x01, x2m3)); + const __nv_bfloat16 y2 = bfloat16_fht_scale(__hadd(x0m1, x23)); + const __nv_bfloat16 y3 = bfloat16_fht_scale(__hsub(x23, x0m1)); + output[base] = y0; + output[base + stride] = y1; + output[base + 2 * stride] = y2; + output[base + 3 * stride] = y3; + const float f0 = fabsf(__bfloat162float(y0)); + const float f1 = fabsf(__bfloat162float(y1)); + const float f2 = fabsf(__bfloat162float(y2)); + const float f3 = fabsf(__bfloat162float(y3)); + return fmaxf(fmaxf(f0, f1), fmaxf(f2, f3)); +} + +template +__device__ __forceinline__ uint8_t quantize_int4_pair_byte( + float x0, + float x1, + float inv_scale, + uint64_t seed, + int64_t idx0) +{ + const int q0 = quantize_int4_value(x0, inv_scale, seed, idx0); + const int q1 = quantize_int4_value(x1, inv_scale, seed, idx0 + 1); + return static_cast(pack_int4_pair(q0, q1)); +} + +template +__device__ __forceinline__ uint32_t quantize_int4_pack4_word( + const T* __restrict__ row_buf, + int col, + float inv_scale, + uint64_t seed, + int64_t row_offset) +{ + const uint32_t b0 = quantize_int4_pair_byte(to_float(row_buf[col]), to_float(row_buf[col + 1]), inv_scale, seed, row_offset + col); + const uint32_t b1 = quantize_int4_pair_byte(to_float(row_buf[col + 2]), to_float(row_buf[col + 3]), inv_scale, seed, row_offset + col + 2); + const uint32_t b2 = quantize_int4_pair_byte(to_float(row_buf[col + 4]), to_float(row_buf[col + 5]), inv_scale, seed, row_offset + col + 4); + const uint32_t b3 = quantize_int4_pair_byte(to_float(row_buf[col + 6]), to_float(row_buf[col + 7]), inv_scale, seed, row_offset + col + 6); + return b0 | (b1 << 8) | (b2 << 16) | (b3 << 24); +} + +template +__device__ __forceinline__ uint32_t quantize_int4_values4_int8_word( + const T* __restrict__ row_buf, + int col, + float inv_scale, + uint64_t seed, + int64_t row_offset) +{ + const uint32_t q0 = static_cast( + quantize_int4_value(to_float(row_buf[col]), inv_scale, seed, row_offset + col)); + const uint32_t q1 = static_cast( + quantize_int4_value(to_float(row_buf[col + 1]), inv_scale, seed, row_offset + col + 1)); + const uint32_t q2 = static_cast( + quantize_int4_value(to_float(row_buf[col + 2]), inv_scale, seed, row_offset + col + 2)); + const uint32_t q3 = static_cast( + quantize_int4_value(to_float(row_buf[col + 3]), inv_scale, seed, row_offset + col + 3)); + return q0 | (q1 << 8) | (q2 << 16) | (q3 << 24); +} + +template +__global__ void quantize_int4_rowwise_kernel( + const InType* __restrict__ x, + int8_t* __restrict__ q, + float* __restrict__ scales, + int K, + uint64_t seed) +{ + constexpr int kWarps = BLOCK_THREADS / kThreadsPerWarp; + __shared__ float warp_smem[kWarps]; + __shared__ float block_smem; + + const int row = static_cast(blockIdx.x); + const int tid = threadIdx.x; + const int64_t row_offset = static_cast(row) * K; + + float absmax = 0.0f; + for (int col = tid; col < K; col += BLOCK_THREADS) { + absmax = fmaxf(absmax, fabsf(to_float(x[row_offset + col]))); + } + absmax = block_reduce_max(absmax, warp_smem, &block_smem); + const float scale = fmaxf(absmax * (1.0f / static_cast(kInt4Max)), 1.0e-10f); + if (tid == 0) { + scales[row] = scale; + } + const float inv_scale = 1.0f / scale; + + const int K_half = K / 2; + for (int packed_col = tid; packed_col < K_half; packed_col += BLOCK_THREADS) { + const int col0 = packed_col * 2; + const int col1 = col0 + 1; + const int q0 = quantize_int4_value(to_float(x[row_offset + col0]), inv_scale, seed, row_offset + col0); + const int q1 = quantize_int4_value(to_float(x[row_offset + col1]), inv_scale, seed, row_offset + col1); + q[static_cast(row) * K_half + packed_col] = pack_int4_pair(q0, q1); + } +} + +template +__global__ void quantize_int4_rowwise_convrot64_kernel( + const InType* __restrict__ x, + int8_t* __restrict__ q, + float* __restrict__ scales, + int K, + uint64_t seed) +{ + constexpr int kGroupThreads = 64; + constexpr int kGroupsInFlight = BLOCK_THREADS / kGroupThreads; + constexpr int kWarps = BLOCK_THREADS / kThreadsPerWarp; + + extern __shared__ __align__(16) unsigned char smem_raw[]; + InType* row_buf = reinterpret_cast(smem_raw); + InType* tmp = row_buf + K; + + __shared__ float warp_smem[kWarps]; + __shared__ float block_smem; + + const int row = static_cast(blockIdx.x); + const int tid = threadIdx.x; + const int sub = tid / kGroupThreads; + const int lane = tid % kGroupThreads; + const int64_t row_offset = static_cast(row) * K; + const int n_groups = K / kConvRotGroup; + + float abs_max = 0.0f; + + if constexpr (PACK4) { + static_assert(BLOCK_THREADS == 640 || BLOCK_THREADS == 768 || BLOCK_THREADS == 960, "PACK4 path is specialized for K=15360"); + #pragma unroll + for (int it = 0; it < 60 / kGroupsInFlight; ++it) { + const int group = it * kGroupsInFlight + sub; + const int base = lane * 4; + const int group_col = group * kConvRotGroup; + const int64_t x_offset = row_offset + group_col + base; + InType* group_buf = row_buf + group_col; + InType* buf0 = tmp + sub * kConvRotGroup; + + convrot_h4_store_typed( + x[x_offset], + x[x_offset + 1], + x[x_offset + 2], + x[x_offset + 3], + group_buf, + base); + __syncwarp(); + + convrot_fht_stage64_vec2_typed(group_buf, buf0, lane, 4); + __syncwarp(); + convrot_fht_stage64_vec2_typed(buf0, group_buf, lane, 16); + __syncthreads(); + + abs_max = fmaxf( + abs_max, + convrot_fht_stage64_store_absmax_typed(group_buf, group_buf, lane, 64)); + } + } else { + InType* buf0 = tmp + sub * (2 * kConvRotGroup); + InType* buf1 = buf0 + kConvRotGroup; + const int iters = (n_groups + kGroupsInFlight - 1) / kGroupsInFlight; + for (int it = 0; it < iters; ++it) { + const int group = it * kGroupsInFlight + sub; + const bool active = group < n_groups; + const int base = lane * 4; + const int group_col = group * kConvRotGroup; + const int64_t x_offset = row_offset + group_col + base; + + const InType x0 = active ? x[x_offset] : from_float(0.0f); + const InType x1 = active ? x[x_offset + 1] : from_float(0.0f); + const InType x2 = active ? x[x_offset + 2] : from_float(0.0f); + const InType x3 = active ? x[x_offset + 3] : from_float(0.0f); + convrot_h4_store_typed(x0, x1, x2, x3, buf1, base); + __syncwarp(); + + convrot_fht_stage64_vec2_typed(buf1, buf0, lane, 4); + __syncwarp(); + convrot_fht_stage64_vec2_typed(buf0, buf1, lane, 16); + __syncthreads(); + + if (active) { + abs_max = fmaxf( + abs_max, + convrot_fht_stage64_store_absmax_typed(buf1, row_buf + group_col, lane, 64)); + } + __syncthreads(); + } + } + + abs_max = block_reduce_max(abs_max, warp_smem, &block_smem); + const float scale = fmaxf(abs_max * (1.0f / static_cast(kInt4Max)), 1.0e-10f); + if (tid == 0) { + scales[row] = scale; + } + const float inv_scale = 1.0f / scale; + + if constexpr (OUTPUT_INT8) { + int8_t* q_row = q + static_cast(row) * K; + uint32_t* q_words = reinterpret_cast(q_row); + const int word_count = K / 4; + for (int word = tid; word < word_count; word += BLOCK_THREADS) { + q_words[word] = quantize_int4_values4_int8_word( + row_buf, word * 4, inv_scale, seed, row_offset); + } + } else if constexpr (PACK4) { + const int K_half = K / 2; + int8_t* q_row = q + static_cast(row) * K_half; + uint32_t* q_words = reinterpret_cast(q_row); + constexpr int kWordCount = 15360 / 8; + for (int word = tid; word < kWordCount; word += BLOCK_THREADS) { + q_words[word] = quantize_int4_pack4_word(row_buf, word * 8, inv_scale, seed, row_offset); + } + } else { + const int K_half = K / 2; + int8_t* q_row = q + static_cast(row) * K_half; + uint32_t* q_words = reinterpret_cast(q_row); + const int word_count = K / 8; + for (int word = tid; word < word_count; word += BLOCK_THREADS) { + q_words[word] = quantize_int4_pack4_word(row_buf, word * 8, inv_scale, seed, row_offset); + } + } +} + +template +__global__ void quantize_int4_rowwise_convrot64_to_int8_float_kernel( + const InType* __restrict__ x, + int8_t* __restrict__ q, + float* __restrict__ scales, + int K, + uint64_t seed) +{ + constexpr int kGroupThreads = 64; + constexpr int kGroupsInFlight = BLOCK_THREADS / kGroupThreads; + constexpr int kWarps = BLOCK_THREADS / kThreadsPerWarp; + + extern __shared__ float smem[]; + float* row_buf = smem; + float* tmp = row_buf + K; + + __shared__ float warp_smem[kWarps]; + __shared__ float block_smem; + + const int row = static_cast(blockIdx.x); + const int tid = threadIdx.x; + const int sub = tid / kGroupThreads; + const int lane = tid % kGroupThreads; + const int64_t row_offset = static_cast(row) * K; + const int n_groups = K / kConvRotGroup; + + float* buf0 = tmp + sub * (2 * kConvRotGroup); + float* buf1 = buf0 + kConvRotGroup; + float abs_max = 0.0f; + + const int iters = (n_groups + kGroupsInFlight - 1) / kGroupsInFlight; + for (int it = 0; it < iters; ++it) { + const int group = it * kGroupsInFlight + sub; + const bool active = group < n_groups; + const int base = lane * 4; + const int group_col = group * kConvRotGroup; + const int64_t x_offset = row_offset + group_col + base; + + const float x0 = active ? to_float(x[x_offset]) : 0.0f; + const float x1 = active ? to_float(x[x_offset + 1]) : 0.0f; + const float x2 = active ? to_float(x[x_offset + 2]) : 0.0f; + const float x3 = active ? to_float(x[x_offset + 3]) : 0.0f; + buf1[base] = 0.5f * (x0 + x1 + x2 - x3); + buf1[base + 1] = 0.5f * (x0 + x1 - x2 + x3); + buf1[base + 2] = 0.5f * (x0 - x1 + x2 + x3); + buf1[base + 3] = 0.5f * (-x0 + x1 + x2 + x3); + __syncthreads(); + + convrot_fht_stage64_float<4>(buf1, buf0, lane); + __syncthreads(); + convrot_fht_stage64_float<16>(buf0, buf1, lane); + __syncthreads(); + + if (active) { + abs_max = fmaxf( + abs_max, + convrot_fht_stage64_store_absmax_float<64>(buf1, row_buf + group_col, lane)); + } + __syncthreads(); + } + + abs_max = block_reduce_max(abs_max, warp_smem, &block_smem); + const float scale = fmaxf(abs_max * (1.0f / static_cast(kInt4Max)), 1.0e-10f); + if (tid == 0) { + scales[row] = scale; + } + const float inv_scale = 1.0f / scale; + + for (int col = tid; col < K; col += BLOCK_THREADS) { + const int64_t idx = row_offset + col; + const float scaled = row_buf[col] * inv_scale; + float quantized; + if constexpr (STOCHASTIC) { + quantized = floorf(scaled + stochastic_rng_value(idx, seed)); + } else { + quantized = nearbyintf(scaled); + } + quantized = fminf(static_cast(kInt4Max), fmaxf(static_cast(-kInt4Max), quantized)); + q[idx] = static_cast(quantized); + } +} + +template +__global__ void quantize_int4_rowwise_convrot_small_kernel( + const InType* __restrict__ x, + int8_t* __restrict__ q, + float* __restrict__ scales, + int K, + uint64_t seed) +{ + static_assert(GROUP_SIZE == 16 || GROUP_SIZE == 64, "small ConvRot int4 kernel supports group sizes 16 and 64"); + constexpr int kGroupThreads = GROUP_SIZE / 4; + constexpr int kGroupsInFlight = BLOCK_THREADS / kGroupThreads; + constexpr int kWarps = BLOCK_THREADS / kThreadsPerWarp; + + extern __shared__ __align__(16) unsigned char smem_raw[]; + InType* row_buf = reinterpret_cast(smem_raw); + InType* tmp = row_buf + K; + + __shared__ float warp_smem[kWarps]; + __shared__ float block_smem; + + const int row = static_cast(blockIdx.x); + const int tid = threadIdx.x; + const int sub = tid / kGroupThreads; + const int lane = tid % kGroupThreads; + const int64_t row_offset = static_cast(row) * K; + const int n_groups = K / GROUP_SIZE; + const int iters = (n_groups + kGroupsInFlight - 1) / kGroupsInFlight; + + InType* buf0 = tmp + sub * (2 * GROUP_SIZE); + InType* buf1 = buf0 + GROUP_SIZE; + float abs_max = 0.0f; + + for (int it = 0; it < iters; ++it) { + const int group = it * kGroupsInFlight + sub; + const bool active = group < n_groups; + const int base = lane * 4; + const int group_col = group * GROUP_SIZE; + const int64_t x_offset = row_offset + group_col + base; + + const InType zero = from_float(0.0f); + const InType x0 = active ? x[x_offset] : zero; + const InType x1 = active ? x[x_offset + 1] : zero; + const InType x2 = active ? x[x_offset + 2] : zero; + const InType x3 = active ? x[x_offset + 3] : zero; + convrot_h4_store_typed(x0, x1, x2, x3, buf1, base); + __syncthreads(); + + if constexpr (GROUP_SIZE == 16) { + if (active) { + abs_max = fmaxf( + abs_max, + convrot_fht_stage64_store_absmax_typed(buf1, row_buf + group_col, lane, 4)); + } + } else { + convrot_fht_stage64_typed(buf1, buf0, lane, 4); + __syncthreads(); + if (active) { + abs_max = fmaxf( + abs_max, + convrot_fht_stage64_store_absmax_typed(buf0, row_buf + group_col, lane, 16)); + } + } + __syncthreads(); + } + + abs_max = block_reduce_max(abs_max, warp_smem, &block_smem); + const float scale = fmaxf(abs_max * (1.0f / static_cast(kInt4Max)), 1.0e-10f); + if (tid == 0) { + scales[row] = scale; + } + const float inv_scale = 1.0f / scale; + + const int K_half = K / 2; + int8_t* q_row = q + static_cast(row) * K_half; + uint32_t* q_words = reinterpret_cast(q_row); + const int word_count = K / 8; + for (int word = tid; word < word_count; word += BLOCK_THREADS) { + q_words[word] = quantize_int4_pack4_word(row_buf, word * 8, inv_scale, seed, row_offset); + } +} + +template +__global__ void dequantize_int4_convrot_small_kernel( + const int8_t* __restrict__ q, + const float* __restrict__ scales, + OutputType* __restrict__ output, + int K, + int scale_size) +{ + static_assert(GROUP_SIZE == 16 || GROUP_SIZE == 64, "small ConvRot int4 dequant kernel supports group sizes 16 and 64"); + constexpr int kGroupThreads = GROUP_SIZE / 4; + extern __shared__ __align__(16) unsigned char smem_raw[]; + OutputType* smem = reinterpret_cast(smem_raw); + + const int sub = threadIdx.x / kGroupThreads; + const int lane = threadIdx.x % kGroupThreads; + const int group = static_cast(blockIdx.x) * GROUPS_PER_BLOCK + sub; + const int row = static_cast(blockIdx.y); + const bool active = !CHECK_BOUNDS || group < K / GROUP_SIZE; + const int64_t row_offset = static_cast(row) * K; + const int group_col = group * GROUP_SIZE; + const float scale = scales[SCALE_PER_ROW ? row : 0]; + + OutputType* buf0 = smem + sub * (2 * GROUP_SIZE); + OutputType* buf1 = buf0 + GROUP_SIZE; + + const int base = lane * 4; + const int64_t q_offset = (row_offset + group_col + base) >> 1; + uint32_t packed0 = 0; + uint32_t packed1 = 0; + if (active) { + packed0 = static_cast(q[q_offset]); + packed1 = static_cast(q[q_offset + 1]); + } + dequant_int4_h4_store_typed( + unpack_int4_nibble(packed0), + unpack_int4_nibble(packed0 >> 4), + unpack_int4_nibble(packed1), + unpack_int4_nibble(packed1 >> 4), + scale, + buf1, + base); + __syncthreads(); + + if constexpr (GROUP_SIZE == 16) { + if (active) { + convrot_fht_stage64_store_typed(buf1, output + row_offset + group_col, lane, 4); + } + } else { + convrot_fht_stage64_typed(buf1, buf0, lane, 4); + __syncthreads(); + if (active) { + convrot_fht_stage64_store_typed(buf0, output + row_offset + group_col, lane, 16); + } + } +} + +template +__global__ void dequantize_int4_convrot64_kernel( + const int8_t* __restrict__ q, + const float* __restrict__ scales, + OutputType* __restrict__ output, + int K, + int scale_size) +{ + constexpr int kGroupThreads = 64; + extern __shared__ __align__(16) unsigned char smem_raw[]; + OutputType* smem = reinterpret_cast(smem_raw); + + const int sub = threadIdx.x / kGroupThreads; + const int lane = threadIdx.x % kGroupThreads; + const int group = static_cast(blockIdx.x) * GROUPS_PER_BLOCK + sub; + const int row = static_cast(blockIdx.y); + const bool active = !CHECK_BOUNDS || group < K / kConvRotGroup; + const int64_t row_offset = static_cast(row) * K; + const int group_col = group * kConvRotGroup; + const int group_col_half = group_col >> 1; + const float scale = scales[SCALE_PER_ROW ? row : 0]; + + OutputType* buf0 = smem + sub * (2 * kConvRotGroup); + OutputType* buf1 = buf0 + kConvRotGroup; + + const int base = lane * 4; + const int64_t q_offset = (static_cast(row) * (K >> 1)) + group_col_half + (base >> 1); + uint32_t packed = 0; + if constexpr (CHECK_BOUNDS) { + if (active) { + packed = static_cast(*reinterpret_cast(q + q_offset)); + } + } else { + packed = static_cast(*reinterpret_cast(q + q_offset)); + } + const int q0 = unpack_int4_nibble(packed); + const int q1 = unpack_int4_nibble(packed >> 4); + const int q2 = unpack_int4_nibble(packed >> 8); + const int q3 = unpack_int4_nibble(packed >> 12); + dequant_int4_h4_store_typed(q0, q1, q2, q3, scale, buf1, base); + __syncwarp(); + + convrot_fht_stage64_vec2_typed(buf1, buf0, lane, 4); + __syncwarp(); + convrot_fht_stage64_vec2_typed(buf0, buf1, lane, 16); + __syncthreads(); + + if constexpr (CHECK_BOUNDS) { + if (active) { + convrot_fht_stage64_store_final_typed(buf1, output + row_offset + group_col, lane); + } + } else { + convrot_fht_stage64_store_final_typed(buf1, output + row_offset + group_col, lane); + } +} + +template +__global__ void dequantize_int4_convrot64_warp32_kernel( + const int8_t* __restrict__ q, + const float* __restrict__ scales, + OutputType* __restrict__ output, + int K, + int scale_size) +{ + constexpr int kGroupThreads = 32; + extern __shared__ __align__(16) unsigned char smem_raw[]; + OutputType* smem = reinterpret_cast(smem_raw); + + const int sub = threadIdx.x / kGroupThreads; + const int lane = threadIdx.x & (kGroupThreads - 1); + const int group = static_cast(blockIdx.x) * GROUPS_PER_BLOCK + sub; + const int row = static_cast(blockIdx.y); + const bool active = !CHECK_BOUNDS || group < K / kConvRotGroup; + const int64_t row_offset = static_cast(row) * K; + const int group_col = group * kConvRotGroup; + const int group_col_half = group_col >> 1; + const float scale = scales[SCALE_PER_ROW ? row : 0]; + + OutputType* buf0 = smem + sub * (2 * kConvRotGroup); + OutputType* buf1 = buf0 + kConvRotGroup; + + const int base = lane * 8; + const int64_t q_offset = (static_cast(row) * (K >> 1)) + group_col_half + (base >> 1); + uint32_t packed = 0; + if constexpr (CHECK_BOUNDS) { + if (active) { + packed = *reinterpret_cast(q + q_offset); + } + } else { + packed = *reinterpret_cast(q + q_offset); + } + + dequant_int4_h4_store_typed( + unpack_int4_nibble(packed), + unpack_int4_nibble(packed >> 4), + unpack_int4_nibble(packed >> 8), + unpack_int4_nibble(packed >> 12), + scale, + buf1, + base); + dequant_int4_h4_store_typed( + unpack_int4_nibble(packed >> 16), + unpack_int4_nibble(packed >> 20), + unpack_int4_nibble(packed >> 24), + unpack_int4_nibble(packed >> 28), + scale, + buf1, + base + 4); + __syncwarp(); + + const int pair_lane = lane * 2; + convrot_fht_stage64_vec2_typed(buf1, buf0, pair_lane, 4); + __syncwarp(); + convrot_fht_stage64_vec2_typed(buf0, buf1, pair_lane, 16); + __syncwarp(); + + if constexpr (CHECK_BOUNDS) { + if (active) { + convrot_fht_stage64_store_final_typed(buf1, output + row_offset + group_col, lane); + } + } else { + convrot_fht_stage64_store_final_typed(buf1, output + row_offset + group_col, lane); + } +} + +constexpr int kStages = 3; +constexpr int kMUnroll = 1; +constexpr int kWarpM = kMUnroll * 16; +constexpr int kNUnroll = 8; +constexpr int kWarpN = kNUnroll * 8; +constexpr int kWarpsM = 2; +constexpr int kWarpsN = 4; +constexpr int kNumWarps = kWarpsM * kWarpsN; +constexpr int kBlockM = kWarpM * kWarpsM; +constexpr int kBlockN = kWarpN * kWarpsN; +constexpr int kBlockKBytes = kGroupSize / 2; +constexpr int kThreadsPerBlock = kNumWarps * 32; +constexpr int kBLoadChunks = kBlockN * 2; +constexpr int kBLoadSweeps = (kBLoadChunks + kThreadsPerBlock - 1) / kThreadsPerBlock; + +template +__global__ void int4_linear_kernel( + const int8_t* __restrict__ act, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scales, + const float* __restrict__ weight_scales, + const BiasType* __restrict__ bias, + OutType* __restrict__ out, + int M, + int N, + int K, + bool has_bias) +{ + const int cta_m = blockIdx.y * kBlockM; + const int cta_n = blockIdx.x * kBlockN; + const int warp_id = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int warp_m = warp_id & (kWarpsM - 1); + const int warp_n = warp_id / kWarpsM; + const int groupID = lane >> 2; + const int tid_in_group = lane & 3; + const int warp_m_base = cta_m + warp_m * kWarpM; + const int warp_n_base = cta_n + warp_n * kWarpN; + + int32_t accum[kMUnroll][kNUnroll][4]; + #pragma unroll + for (int mi = 0; mi < kMUnroll; ++mi) { + #pragma unroll + for (int c = 0; c < kNUnroll; ++c) { + #pragma unroll + for (int i = 0; i < 4; ++i) { + accum[mi][c][i] = 0; + } + } + } + + const int K_half = K / 2; + const int num_groups = K / kGroupSize; + + __shared__ alignas(16) int8_t smem_B[kStages][kBlockN * kBlockKBytes]; + __shared__ alignas(16) int8_t smem_A[kStages][kBlockM * kBlockKBytes]; + + auto issue_B_load = [&](int g, int stage) { + if (g >= num_groups) return; + const int thread_idx = threadIdx.x; + #pragma unroll + for (int sweep = 0; sweep < kBLoadSweeps; ++sweep) { + const int t = thread_idx + sweep * kThreadsPerBlock; + if (t < kBlockN * 2) { + const int n_row = t >> 1; + const int half = t & 1; + const int n_global = cta_n + n_row; + int8_t* dst = &smem_B[stage][n_row * kBlockKBytes + half * 16]; + if (n_global < N) { + const int8_t* src = weight + n_global * K_half + g * kBlockKBytes + half * 16; + cp_async_16b(dst, src); + } else { + reinterpret_cast(dst)[0] = {0, 0, 0, 0}; + } + } + } + }; + + auto issue_A_load = [&](int g, int stage) { + if (g >= num_groups) return; + const int t = threadIdx.x; + if (t < kBlockM * 2) { + const int m_row = t >> 1; + const int half = t & 1; + const int m_global = cta_m + m_row; + int8_t* dst = &smem_A[stage][m_row * kBlockKBytes + half * 16]; + if (m_global < M) { + const int8_t* src = act + m_global * K_half + g * kBlockKBytes + half * 16; + cp_async_16b(dst, src); + } else { + reinterpret_cast(dst)[0] = {0, 0, 0, 0}; + } + } + }; + + auto load_B_fragment = [&](int stage, int c, uint32_t (&b_reg)[2]) { + const int b_col_local = (warp_n * kWarpN) + c * 8 + groupID; + const int b_col_global = cta_n + b_col_local; + b_reg[0] = b_reg[1] = 0; + if (b_col_local < kBlockN && b_col_global < N) { + const int byte0 = tid_in_group * 8; + const int8_t* row_base = &smem_B[stage][b_col_local * kBlockKBytes]; + b_reg[0] = *reinterpret_cast(row_base + byte0); + b_reg[1] = *reinterpret_cast(row_base + byte0 + 4); + } + }; + + #pragma unroll + for (int s = 0; s < kStages - 1; ++s) { + issue_A_load(s, s); + issue_B_load(s, s); + cp_async_commit_group(); + } + + for (int g = 0; g < num_groups; ++g) { + const int next_g = g + kStages - 1; + if (next_g < num_groups) { + const int next_stage = (g + kStages - 1) % kStages; + issue_A_load(next_g, next_stage); + issue_B_load(next_g, next_stage); + } + cp_async_commit_group(); + cp_async_wait_group(); + __syncthreads(); + + const int cur_stage = g % kStages; + uint32_t a_reg[kMUnroll][4]; + #pragma unroll + for (int mi = 0; mi < kMUnroll; ++mi) { + const int m_tile_base = warp_m_base + mi * 16; + const int row0_m = m_tile_base + groupID; + const int row1_m = m_tile_base + groupID + 8; + const int row0_local = warp_m * kWarpM + mi * 16 + groupID; + const int row1_local = warp_m * kWarpM + mi * 16 + groupID + 8; + a_reg[mi][0] = a_reg[mi][1] = a_reg[mi][2] = a_reg[mi][3] = 0; + if (row0_m < M) { + const int8_t* rb = &smem_A[cur_stage][row0_local * kBlockKBytes]; + a_reg[mi][0] = *reinterpret_cast(rb + tid_in_group * 8); + a_reg[mi][2] = *reinterpret_cast(rb + tid_in_group * 8 + 4); + } + if (row1_m < M) { + const int8_t* rb = &smem_A[cur_stage][row1_local * kBlockKBytes]; + a_reg[mi][1] = *reinterpret_cast(rb + tid_in_group * 8); + a_reg[mi][3] = *reinterpret_cast(rb + tid_in_group * 8 + 4); + } + } + + #pragma unroll + for (int c = 0; c < kNUnroll; ++c) { + uint32_t b_reg[2]; + load_B_fragment(cur_stage, c, b_reg); + #pragma unroll + for (int mi = 0; mi < kMUnroll; ++mi) { + int32_t zero[4] = {0, 0, 0, 0}; + int32_t d[4]; + mma_m16n8k64_s4s4s32(a_reg[mi], b_reg, zero, d); + accum[mi][c][0] += d[0]; + accum[mi][c][1] += d[1]; + accum[mi][c][2] += d[2]; + accum[mi][c][3] += d[3]; + } + } + } + cp_async_wait_group<0>(); + + #pragma unroll + for (int mi = 0; mi < kMUnroll; ++mi) { + const int m_tile_base = warp_m_base + mi * 16; + const int row0_m = m_tile_base + groupID; + const int row1_m = m_tile_base + groupID + 8; + #pragma unroll + for (int c = 0; c < kNUnroll; ++c) { + const int n_chunk_base = warp_n_base + c * 8; + const int col0 = n_chunk_base + tid_in_group * 2 + 0; + const int col1 = n_chunk_base + tid_in_group * 2 + 1; + if (row0_m < M && col0 < N) { + float v = static_cast(accum[mi][c][0]) * x_scales[row0_m] * weight_scales[col0]; + if (has_bias) v += to_float(bias[col0]); + out[row0_m * N + col0] = from_float(v); + } + if (row0_m < M && col1 < N) { + float v = static_cast(accum[mi][c][1]) * x_scales[row0_m] * weight_scales[col1]; + if (has_bias) v += to_float(bias[col1]); + out[row0_m * N + col1] = from_float(v); + } + if (row1_m < M && col0 < N) { + float v = static_cast(accum[mi][c][2]) * x_scales[row1_m] * weight_scales[col0]; + if (has_bias) v += to_float(bias[col0]); + out[row1_m * N + col0] = from_float(v); + } + if (row1_m < M && col1 < N) { + float v = static_cast(accum[mi][c][3]) * x_scales[row1_m] * weight_scales[col1]; + if (has_bias) v += to_float(bias[col1]); + out[row1_m * N + col1] = from_float(v); + } + } + } +} + +__global__ void unpack_int4_to_int8_kernel( + const int8_t* __restrict__ input, + int8_t* __restrict__ output, + int64_t total_packed, + int K_half) +{ + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total_packed) { + return; + } + const int row = static_cast(idx / K_half); + const int packed_col = static_cast(idx - static_cast(row) * K_half); + const uint32_t packed = static_cast(input[idx]); + int8_t* out_row = output + static_cast(row) * (K_half * 2); + out_row[packed_col * 2] = static_cast(unpack_int4_nibble(packed)); + out_row[packed_col * 2 + 1] = static_cast(unpack_int4_nibble(packed >> 4)); +} + +} // namespace + +extern "C" { + +void launch_quantize_int4_rowwise_kernel( + const void* input, + void* output, + void* scales, + int64_t M, + int64_t K, + int input_dtype_code, + bool stochastic, + uint64_t seed, + cudaStream_t stream) +{ + if (K % comfy::svdquant::kGroupSize != 0) return; + constexpr int kThreads = 256; + const dim3 grid(static_cast(M)); + const dim3 block(kThreads); + if (input_dtype_code == 2) { + if (stochastic) { + quantize_int4_rowwise_kernel<__nv_bfloat16, kThreads, true> + <<>>(reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(scales), + static_cast(K), + seed); + } else { + quantize_int4_rowwise_kernel<__nv_bfloat16, kThreads, false> + <<>>(reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(scales), + static_cast(K), + seed); + } + } else if (input_dtype_code == 1) { + if (stochastic) { + quantize_int4_rowwise_kernel<__half, kThreads, true> + <<>>(reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(scales), + static_cast(K), + seed); + } else { + quantize_int4_rowwise_kernel<__half, kThreads, false> + <<>>(reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(scales), + static_cast(K), + seed); + } + } else if (input_dtype_code == 0) { + if (stochastic) { + quantize_int4_rowwise_kernel + <<>>(reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(scales), + static_cast(K), + seed); + } else { + quantize_int4_rowwise_kernel + <<>>(reinterpret_cast(input), + reinterpret_cast(output), + reinterpret_cast(scales), + static_cast(K), + seed); + } + } +} + +void launch_quantize_int4_rowwise_convrot64_kernel( + const void* input, + void* output, + void* scales, + int64_t M, + int64_t K, + int group_size, + int input_dtype_code, + bool stochastic, + uint64_t seed, + cudaStream_t stream) +{ + if ((group_size != 16 && group_size != 64 && group_size != kConvRotGroup) || K % group_size != 0) return; + constexpr int block_threads_multi = 1024; + constexpr int block_threads_15360 = 640; + constexpr int block_threads_single = 512; + constexpr int block_threads_small = 256; + + auto launch = [&](auto kernel, auto typed_input, int block_threads, int scratch_buffers) { + using TypedInputPtr = decltype(typed_input); + using TypedInput = std::remove_cv_t>; + const int groups_in_flight = block_threads / 64; + const size_t smem_bytes = + (static_cast(K) + groups_in_flight * scratch_buffers * kConvRotGroup) * sizeof(TypedInput); + cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_bytes)); + kernel<<(M), block_threads, smem_bytes, stream>>>( + typed_input, + reinterpret_cast(output), + reinterpret_cast(scales), + static_cast(K), + seed); + }; + auto launch_small = [&](auto kernel, auto typed_input, int small_group_size, int block_threads_small_group) { + using TypedInputPtr = decltype(typed_input); + using TypedInput = std::remove_cv_t>; + const int groups_in_flight = block_threads_small_group / (small_group_size / 4); + const size_t smem_bytes = + (static_cast(K) + groups_in_flight * 2 * small_group_size) * sizeof(TypedInput); + cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_bytes)); + kernel<<(M), block_threads_small_group, smem_bytes, stream>>>( + typed_input, + reinterpret_cast(output), + reinterpret_cast(scales), + static_cast(K), + seed); + }; + if (group_size == 16) { + const bool use_wide_small_group = K > 4096; + if (input_dtype_code == 2) { + auto typed_input = reinterpret_cast(input); + if (use_wide_small_group) { + if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, __nv_bfloat16, 512, true>, typed_input, 16, 512); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, __nv_bfloat16, 512, false>, typed_input, 16, 512); + } + } else if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, __nv_bfloat16, 128, true>, typed_input, 16, 128); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, __nv_bfloat16, 128, false>, typed_input, 16, 128); + } + } else if (input_dtype_code == 1) { + auto typed_input = reinterpret_cast(input); + if (use_wide_small_group) { + if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, __half, 512, true>, typed_input, 16, 512); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, __half, 512, false>, typed_input, 16, 512); + } + } else if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, __half, 128, true>, typed_input, 16, 128); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, __half, 128, false>, typed_input, 16, 128); + } + } else if (input_dtype_code == 0) { + auto typed_input = reinterpret_cast(input); + if (use_wide_small_group) { + if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, float, 512, true>, typed_input, 16, 512); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, float, 512, false>, typed_input, 16, 512); + } + } else if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, float, 128, true>, typed_input, 16, 128); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<16, float, 128, false>, typed_input, 16, 128); + } + } + return; + } + if (group_size == 64) { + const bool use_wide_small_group = K > 4096; + if (input_dtype_code == 2) { + auto typed_input = reinterpret_cast(input); + if (use_wide_small_group) { + if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, __nv_bfloat16, 512, true>, typed_input, 64, 512); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, __nv_bfloat16, 512, false>, typed_input, 64, 512); + } + } else if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, __nv_bfloat16, 128, true>, typed_input, 64, 128); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, __nv_bfloat16, 128, false>, typed_input, 64, 128); + } + } else if (input_dtype_code == 1) { + auto typed_input = reinterpret_cast(input); + if (use_wide_small_group) { + if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, __half, 512, true>, typed_input, 64, 512); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, __half, 512, false>, typed_input, 64, 512); + } + } else if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, __half, 128, true>, typed_input, 64, 128); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, __half, 128, false>, typed_input, 64, 128); + } + } else if (input_dtype_code == 0) { + auto typed_input = reinterpret_cast(input); + if (use_wide_small_group) { + if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, float, 512, true>, typed_input, 64, 512); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, float, 512, false>, typed_input, 64, 512); + } + } else if (stochastic) { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, float, 128, true>, typed_input, 64, 128); + } else { + launch_small(quantize_int4_rowwise_convrot_small_kernel<64, float, 128, false>, typed_input, 64, 128); + } + } + return; + } + if (input_dtype_code == 2) { + auto typed_input = reinterpret_cast(input); + if (M != 1 && K <= 4096) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_small, false, true>, typed_input, block_threads_small, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_small, false, false>, typed_input, block_threads_small, 2); + } + } else if (M == 1) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_single, false, true>, typed_input, block_threads_single, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_single, false, false>, typed_input, block_threads_single, 2); + } + } else if (K == 15360) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_15360, true, true>, typed_input, block_threads_15360, 1); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_15360, true, false>, typed_input, block_threads_15360, 1); + } + } else { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_multi, false, true>, typed_input, block_threads_multi, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_multi, false, false>, typed_input, block_threads_multi, 2); + } + } + } else if (input_dtype_code == 1) { + auto typed_input = reinterpret_cast(input); + if (M != 1 && K <= 4096) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_small, false, true>, typed_input, block_threads_small, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_small, false, false>, typed_input, block_threads_small, 2); + } + } else if (M == 1) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_single, false, true>, typed_input, block_threads_single, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_single, false, false>, typed_input, block_threads_single, 2); + } + } else if (K == 15360) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_15360, true, true>, typed_input, block_threads_15360, 1); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_15360, true, false>, typed_input, block_threads_15360, 1); + } + } else { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_multi, false, true>, typed_input, block_threads_multi, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_multi, false, false>, typed_input, block_threads_multi, 2); + } + } + } else if (input_dtype_code == 0) { + auto typed_input = reinterpret_cast(input); + if (M != 1 && K <= 4096) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_small, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_small, 2); + } + } else if (M == 1) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_single, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_single, 2); + } + } else if (K == 15360) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_15360, 1); + } else { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_15360, 1); + } + } else { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_multi, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_multi, 2); + } + } + } +} + +void launch_quantize_int4_rowwise_convrot64_to_int8_kernel( + const void* input, + void* output, + void* scales, + int64_t M, + int64_t K, + int group_size, + int input_dtype_code, + bool stochastic, + uint64_t seed, + cudaStream_t stream) +{ + if (group_size != kConvRotGroup || K % kConvRotGroup != 0) return; + constexpr int block_threads_multi = 1024; + constexpr int block_threads_15360 = 640; + constexpr int block_threads_single = 512; + + auto launch = [&](auto kernel, auto typed_input, int block_threads, int scratch_buffers) { + using TypedInputPtr = decltype(typed_input); + using TypedInput = std::remove_cv_t>; + const int groups_in_flight = block_threads / 64; + const size_t smem_bytes = + (static_cast(K) + groups_in_flight * scratch_buffers * kConvRotGroup) * sizeof(TypedInput); + cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_bytes)); + kernel<<(M), block_threads, smem_bytes, stream>>>( + typed_input, + reinterpret_cast(output), + reinterpret_cast(scales), + static_cast(K), + seed); + }; + + auto launch_float = [&](auto kernel, auto typed_input, int block_threads) { + const int groups_in_flight = block_threads / 64; + const size_t smem_bytes = + (static_cast(K) + groups_in_flight * 2 * kConvRotGroup) * sizeof(float); + cudaFuncSetAttribute( + kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(smem_bytes)); + kernel<<(M), block_threads, smem_bytes, stream>>>( + typed_input, + reinterpret_cast(output), + reinterpret_cast(scales), + static_cast(K), + seed); + }; + + if (input_dtype_code == 2) { + auto typed_input = reinterpret_cast(input); + if (K <= 4096) { + if (stochastic) { + launch_float( + quantize_int4_rowwise_convrot64_to_int8_float_kernel<__nv_bfloat16, block_threads_multi, true>, + typed_input, + block_threads_multi); + } else { + launch_float( + quantize_int4_rowwise_convrot64_to_int8_float_kernel<__nv_bfloat16, block_threads_multi, false>, + typed_input, + block_threads_multi); + } + } else if (M == 1) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_single, false, true, true>, typed_input, block_threads_single, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_single, false, false, true>, typed_input, block_threads_single, 2); + } + } else if (K == 15360) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_15360, true, true, true>, typed_input, block_threads_15360, 1); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_15360, true, false, true>, typed_input, block_threads_15360, 1); + } + } else { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_multi, false, true, true>, typed_input, block_threads_multi, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__nv_bfloat16, block_threads_multi, false, false, true>, typed_input, block_threads_multi, 2); + } + } + } else if (input_dtype_code == 1) { + auto typed_input = reinterpret_cast(input); + if (K <= 4096) { + if (stochastic) { + launch_float( + quantize_int4_rowwise_convrot64_to_int8_float_kernel<__half, block_threads_multi, true>, + typed_input, + block_threads_multi); + } else { + launch_float( + quantize_int4_rowwise_convrot64_to_int8_float_kernel<__half, block_threads_multi, false>, + typed_input, + block_threads_multi); + } + } else if (M == 1) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_single, false, true, true>, typed_input, block_threads_single, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_single, false, false, true>, typed_input, block_threads_single, 2); + } + } else if (K == 15360) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_15360, true, true, true>, typed_input, block_threads_15360, 1); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_15360, true, false, true>, typed_input, block_threads_15360, 1); + } + } else { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_multi, false, true, true>, typed_input, block_threads_multi, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel<__half, block_threads_multi, false, false, true>, typed_input, block_threads_multi, 2); + } + } + } else if (input_dtype_code == 0) { + auto typed_input = reinterpret_cast(input); + if (K <= 4096) { + if (stochastic) { + launch_float( + quantize_int4_rowwise_convrot64_to_int8_float_kernel, + typed_input, + block_threads_multi); + } else { + launch_float( + quantize_int4_rowwise_convrot64_to_int8_float_kernel, + typed_input, + block_threads_multi); + } + } else if (M == 1) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_single, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_single, 2); + } + } else if (K == 15360) { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_15360, 1); + } else { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_15360, 1); + } + } else { + if (stochastic) { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_multi, 2); + } else { + launch(quantize_int4_rowwise_convrot64_kernel, typed_input, block_threads_multi, 2); + } + } + } +} + +void launch_dequantize_int4_convrot64_kernel( + const void* input, + const void* scales, + void* output, + int64_t num_rows, + int64_t num_cols, + int64_t scale_size, + int group_size, + int output_dtype_code, + cudaStream_t stream) +{ + if (num_rows == 0 || num_cols == 0) return; + if ((group_size != 16 && group_size != 64 && group_size != kConvRotGroup) || num_cols % group_size != 0) return; + + auto launch_small_groups = [&](auto group_tag, auto groups_tag) { + constexpr int small_group_size = decltype(group_tag)::value; + constexpr int groups_per_block = decltype(groups_tag)::value; + constexpr int block_threads = groups_per_block * (small_group_size / 4); + const int group_blocks = + static_cast((num_cols / small_group_size + groups_per_block - 1) / groups_per_block); + dim3 grid(static_cast(group_blocks), static_cast(num_rows)); + const bool check_bounds = ((num_cols / small_group_size) % groups_per_block) != 0; + const bool scale_per_row = scale_size != 1; + + if (output_dtype_code == 2) { + const size_t smem_bytes = groups_per_block * 2 * small_group_size * sizeof(__nv_bfloat16); + if (check_bounds) { + if (scale_per_row) { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } else { + if (scale_per_row) { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } + } else if (output_dtype_code == 1) { + const size_t smem_bytes = groups_per_block * 2 * small_group_size * sizeof(__half); + if (check_bounds) { + if (scale_per_row) { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } else { + if (scale_per_row) { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } + } else { + const size_t smem_bytes = groups_per_block * 2 * small_group_size * sizeof(float); + if (check_bounds) { + if (scale_per_row) { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } else { + if (scale_per_row) { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot_small_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } + } + }; + + if (group_size == 16) { + launch_small_groups(std::integral_constant{}, std::integral_constant{}); + return; + } + if (group_size == 64) { + launch_small_groups(std::integral_constant{}, std::integral_constant{}); + return; + } + + auto launch_groups = [&](auto groups_tag) { + constexpr int groups_per_block = decltype(groups_tag)::value; + constexpr int block_threads = groups_per_block * 64; + const int group_blocks = + static_cast((num_cols / kConvRotGroup + groups_per_block - 1) / groups_per_block); + dim3 grid(static_cast(group_blocks), static_cast(num_rows)); + const bool check_bounds = ((num_cols / kConvRotGroup) % groups_per_block) != 0; + const bool scale_per_row = scale_size != 1; + + if (output_dtype_code == 2) { + const size_t smem_bytes = groups_per_block * 2 * kConvRotGroup * sizeof(__nv_bfloat16); + if (check_bounds) { + if (scale_per_row) { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } else { + if (scale_per_row) { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } + } else if (output_dtype_code == 1) { + const size_t smem_bytes = groups_per_block * 2 * kConvRotGroup * sizeof(__half); + if (check_bounds) { + if (scale_per_row) { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } else { + if (scale_per_row) { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } + } else { + const size_t smem_bytes = groups_per_block * 2 * kConvRotGroup * sizeof(float); + if (check_bounds) { + if (scale_per_row) { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } else { + if (scale_per_row) { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot64_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } + } + }; + + auto launch_groups_warp32 = [&](auto groups_tag) { + constexpr int groups_per_block = decltype(groups_tag)::value; + constexpr int block_threads = groups_per_block * 32; + const int group_blocks = + static_cast((num_cols / kConvRotGroup + groups_per_block - 1) / groups_per_block); + dim3 grid(static_cast(group_blocks), static_cast(num_rows)); + const bool check_bounds = ((num_cols / kConvRotGroup) % groups_per_block) != 0; + const bool scale_per_row = scale_size != 1; + + if (output_dtype_code == 2) { + const size_t smem_bytes = groups_per_block * 2 * kConvRotGroup * sizeof(__nv_bfloat16); + if (check_bounds) { + if (scale_per_row) { + dequantize_int4_convrot64_warp32_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot64_warp32_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } else { + if (scale_per_row) { + dequantize_int4_convrot64_warp32_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot64_warp32_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__nv_bfloat16*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } + } else { + const size_t smem_bytes = groups_per_block * 2 * kConvRotGroup * sizeof(__half); + if (check_bounds) { + if (scale_per_row) { + dequantize_int4_convrot64_warp32_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot64_warp32_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } else { + if (scale_per_row) { + dequantize_int4_convrot64_warp32_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } else { + dequantize_int4_convrot64_warp32_kernel + <<>>( + reinterpret_cast(input), + reinterpret_cast(scales), + reinterpret_cast<__half*>(output), + static_cast(num_cols), + static_cast(scale_size)); + } + } + } + }; + + if (output_dtype_code == 1 || output_dtype_code == 2) { + if (num_cols < 1024) { + launch_groups_warp32(std::integral_constant{}); + } else if (num_cols < 4096) { + launch_groups_warp32(std::integral_constant{}); + } else { + launch_groups_warp32(std::integral_constant{}); + } + } else { + if (num_cols < 1024) { + launch_groups(std::integral_constant{}); + } else if (num_cols < 4096) { + launch_groups(std::integral_constant{}); + } else if (num_cols < 8192) { + launch_groups(std::integral_constant{}); + } else { + launch_groups(std::integral_constant{}); + } + } +} + +void launch_int4_linear_kernel( + const void* act, + const void* weight, + const void* x_scales, + const void* weight_scales, + const void* bias, + void* output, + int64_t M, + int64_t N, + int64_t K, + bool has_bias, + int output_dtype_code, + int bias_dtype_code, + cudaStream_t stream) +{ + if (K % comfy::svdquant::kGroupSize != 0) return; + const dim3 grid(static_cast((N + kBlockN - 1) / kBlockN), + static_cast((M + kBlockM - 1) / kBlockM)); + const dim3 block(kThreadsPerBlock); + +#define DISPATCH_OUT_BIAS(OutType, BiasType) \ + int4_linear_kernel<<>>( \ + reinterpret_cast(act), \ + reinterpret_cast(weight), \ + reinterpret_cast(x_scales), \ + reinterpret_cast(weight_scales), \ + reinterpret_cast(bias), \ + reinterpret_cast(output), \ + static_cast(M), static_cast(N), static_cast(K), has_bias) + +#define DISPATCH_BIAS(OutType) \ + do { \ + if (bias_dtype_code == 2) DISPATCH_OUT_BIAS(OutType, __nv_bfloat16); \ + else if (bias_dtype_code == 1) DISPATCH_OUT_BIAS(OutType, __half); \ + else DISPATCH_OUT_BIAS(OutType, float); \ + } while (0) + + if (output_dtype_code == 2) { + DISPATCH_BIAS(__nv_bfloat16); + } else if (output_dtype_code == 1) { + DISPATCH_BIAS(__half); + } else { + DISPATCH_BIAS(float); + } + +#undef DISPATCH_BIAS +#undef DISPATCH_OUT_BIAS +} + +void launch_unpack_int4_to_int8_kernel( + const void* input, + void* output, + int64_t rows, + int64_t K_half, + cudaStream_t stream) +{ + if (rows == 0 || K_half == 0) return; + const int64_t total_packed = rows * K_half; + constexpr int block_threads = 256; + const int64_t blocks = (total_packed + block_threads - 1) / block_threads; + unpack_int4_to_int8_kernel<<(blocks), block_threads, 0, stream>>>( + reinterpret_cast(input), + reinterpret_cast(output), + total_packed, + static_cast(K_half)); +} + +bool launch_cutlass_int4_dequant( + const void* A, + const void* B, + const void* xs, + const void* ws, + const void* bias, + void* D, + int64_t M, + int64_t N, + int64_t K, + int out_dtype_code, + cudaStream_t stream) +{ +#ifdef COMFY_HAVE_CUTLASS + if (M == 0 || N == 0 || K == 0) return true; + if (K % comfy::svdquant::kGroupSize != 0) return false; + const int8_t* a = static_cast(A); + const int8_t* b = static_cast(B); + const float* x = static_cast(xs); + const float* w = static_cast(ws); + const float* bs = static_cast(bias); + if (bs == nullptr) { + switch (out_dtype_code) { + case 0: return dispatch_fused_int4_no_bias(a, b, x, w, static_cast(D), M, N, K, stream); + case 1: return dispatch_fused_int4_no_bias(a, b, x, w, static_cast(D), M, N, K, stream); + case 2: return dispatch_fused_int4_no_bias(a, b, x, w, static_cast(D), M, N, K, stream); + default: return false; + } + } + switch (out_dtype_code) { + case 0: return dispatch_fused_int4(a, b, x, w, bs, static_cast(D), M, N, K, stream); + case 1: return dispatch_fused_int4(a, b, x, w, bs, static_cast(D), M, N, K, stream); + case 2: return dispatch_fused_int4(a, b, x, w, bs, static_cast(D), M, N, K, stream); + default: return false; + } +#else + return false; +#endif +} + +} // extern "C" diff --git a/comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu b/comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu index 5bb6956..b526b96 100644 --- a/comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu +++ b/comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu @@ -275,7 +275,7 @@ bool launch_cutlass_int8_dequant( const float* x = static_cast(xs); const float* w = static_cast(ws); const float* bs = static_cast(bias); - if (bs == nullptr && M <= 1024) { + if (bs == nullptr) { switch (out_dtype_code) { case 0: return dispatch_fused_no_bias(a, b, x, w, static_cast(D), M, N, K, stream); case 1: return dispatch_fused_no_bias(a, b, x, w, static_cast(D), M, N, K, stream); diff --git a/comfy_kitchen/backends/eager/__init__.py b/comfy_kitchen/backends/eager/__init__.py index 0200296..0850ee0 100644 --- a/comfy_kitchen/backends/eager/__init__.py +++ b/comfy_kitchen/backends/eager/__init__.py @@ -11,10 +11,14 @@ "dequantize_int8_simple_dtype", "dequantize_int8_convrot_weight", "dequantize_int8_convrot_weight_dtype", + "dequantize_convrot_w4a4_weight", "gemv_awq_w4a16", + "convrot_w4a4_linear", + "prepare_int4_weight_for_int8_linear", "quantize_mxfp8", "quantize_nvfp4", "quantize_per_tensor_fp8", + "quantize_convrot_w4a4_weight", "quantize_int8_rowwise", "quantize_int8_convrot_weight", "quantize_and_rotate_rowwise", @@ -38,6 +42,12 @@ from .adaln import adaln from .awq import gemv_awq_w4a16 +from .convrot_w4a4 import ( + convrot_w4a4_linear, + dequantize_convrot_w4a4_weight, + prepare_int4_weight_for_int8_linear, + quantize_convrot_w4a4_weight, +) from .quantization import ( dequantize_int8_convrot_weight, dequantize_int8_convrot_weight_dtype, @@ -211,6 +221,42 @@ def _build_constraints() -> dict: }, default_devices=all_devices, ), + "quantize_convrot_w4a4_weight": FunctionConstraints( + params={ + "weight": ParamConstraint(dtypes=standard_floats, shape_rules=(ExactDims(2),)), + "convrot_groupsize": ParamConstraint(dtypes=frozenset({int})), + "quant_group_size": ParamConstraint(dtypes=frozenset({int})), + "stochastic_rounding": ParamConstraint(dtypes=frozenset({int})), + }, + default_devices=all_devices, + ), + "dequantize_convrot_w4a4_weight": FunctionConstraints( + params={ + "qdata": ParamConstraint(dtypes=frozenset({torch.int8}), shape_rules=(ExactDims(2),)), + "scales": ParamConstraint(dtypes=standard_floats, shape_rules=(ExactDims(1),)), + "convrot_groupsize": ParamConstraint(dtypes=frozenset({int})), + "quant_group_size": ParamConstraint(dtypes=frozenset({int})), + "output_dtype": ParamConstraint(dtypes=standard_floats), + }, + default_devices=all_devices, + ), + "convrot_w4a4_linear": FunctionConstraints( + params={ + "x": ParamConstraint(dtypes=standard_floats), + "qweight": ParamConstraint(dtypes=frozenset({torch.int8}), shape_rules=(ExactDims(2),)), + "wscales": ParamConstraint(dtypes=standard_floats, shape_rules=(ExactDims(1),)), + "bias": ParamConstraint(dtypes=standard_floats), + "convrot_groupsize": ParamConstraint(dtypes=frozenset({int})), + "quant_group_size": ParamConstraint(dtypes=frozenset({int})), + }, + default_devices=all_devices, + ), + "prepare_int4_weight_for_int8_linear": FunctionConstraints( + params={ + "weight": ParamConstraint(dtypes=frozenset({torch.int8}), shape_rules=(ExactDims(2),)), + }, + default_devices=all_devices, + ), "gemv_awq_w4a16": FunctionConstraints( params={ "x": ParamConstraint(dtypes=standard_floats), diff --git a/comfy_kitchen/backends/eager/convrot_w4a4.py b/comfy_kitchen/backends/eager/convrot_w4a4.py new file mode 100644 index 0000000..0e08155 --- /dev/null +++ b/comfy_kitchen/backends/eager/convrot_w4a4.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 Comfy Org. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Eager ConvRot W4A4 kernels.""" +from __future__ import annotations + +import math + +import torch + +from .svdquant import _INT4_GROUP_SIZE, _INT4_MAX, _pack_int4_row_major, _unpack_int4_row_major + +_HADAMARD_CACHE = {} + + +def _build_hadamard( + size: int, + device: str | torch.device = "cpu", + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + cache_key = (size, str(device), dtype) + if cache_key in _HADAMARD_CACHE: + return _HADAMARD_CACHE[cache_key] + + if size < 4 or (size & (size - 1)) != 0 or math.log(size, 4) % 1 != 0: + raise ValueError(f"Regular Hadamard size must be a power of 4, got {size}") + + h4 = torch.tensor( + [[1, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]], + dtype=dtype, + device=device, + ) + + h = h4 + current_size = 4 + while current_size < size: + h = torch.kron(h, h4) + current_size *= 4 + + h_normalized = h / (size**0.5) + _HADAMARD_CACHE[cache_key] = h_normalized + return h_normalized + + +def _rotate_weight( + weight: torch.Tensor, + h: torch.Tensor, + group_size: int, +) -> torch.Tensor: + out_f, in_f = weight.shape + if in_f % group_size != 0: + raise ValueError(f"in_features {in_f} not divisible by group_size {group_size}") + n_groups = in_f // group_size + + weight_grouped = weight.reshape(out_f, n_groups, group_size) + h_t = h.T.to(dtype=weight.dtype, device=weight.device) + weight_rotated = torch.matmul(weight_grouped, h_t) + return weight_rotated.reshape(out_f, in_f) + + +def _rotate_activation( + x: torch.Tensor, + h: torch.Tensor, + group_size: int, +) -> torch.Tensor: + orig_shape = x.shape + features = orig_shape[-1] + if features % group_size != 0: + raise ValueError(f"features {features} not divisible by group_size {group_size}") + n_groups = features // group_size + + x_grouped = x.reshape(-1, n_groups, group_size) + h = h.to(dtype=x.dtype, device=x.device) + x_rotated = torch.matmul(x_grouped, h) + return x_rotated.reshape(orig_shape) + + +def validate_w4a4_shape(tensor: torch.Tensor, convrot_groupsize: int, quant_group_size: int) -> None: + if tensor.dim() != 2: + raise ValueError(f"ConvRot W4A4 expects a 2D tensor, got shape {tuple(tensor.shape)}") + k = tensor.shape[-1] + if k % convrot_groupsize != 0: + raise ValueError(f"in_features {k} not divisible by convrot_groupsize {convrot_groupsize}") + if k % quant_group_size != 0: + raise ValueError(f"in_features {k} not divisible by quant_group_size {quant_group_size}") + + +def _int4_stochastic_rng(x: torch.Tensor, seed: int) -> torch.Tensor: + generator = torch.Generator(device=x.device) + generator.manual_seed(seed) + return torch.rand( + x.shape, + dtype=x.dtype, + layout=x.layout, + device=x.device, + generator=generator, + ) + + +def _round_int4(scaled: torch.Tensor, stochastic_rounding: int | None = 0) -> torch.Tensor: + if stochastic_rounding is not None and stochastic_rounding > 0: + rng = _int4_stochastic_rng(scaled, stochastic_rounding) + scaled.add_(rng) + return scaled.floor_().clamp_(-_INT4_MAX, _INT4_MAX).to(torch.int8) + return scaled.round_().clamp_(-_INT4_MAX, _INT4_MAX).to(torch.int8) + + +def quantize_signed_int4_rowwise( + x: torch.Tensor, + stochastic_rounding: int | None = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + rows, _ = x.shape + absmax = x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-10) + scales = absmax / _INT4_MAX + q = _round_int4(x / scales, stochastic_rounding=stochastic_rounding) + return _pack_int4_row_major(q), scales.reshape(rows).to(torch.float32) + + +def prepare_int4_weight_for_int8_linear(weight: torch.Tensor) -> torch.Tensor: + """Unpack a packed signed INT4 weight matrix into INT8 values.""" + if weight.dim() != 2: + raise ValueError("prepared INT4 fallback weight expects a 2D tensor") + return _unpack_int4_row_major(weight).to(torch.int8) + + +def quantize_convrot_w4a4_weight( + weight: torch.Tensor, + convrot_groupsize: int = 256, + quant_group_size: int = _INT4_GROUP_SIZE, + stochastic_rounding: int | None = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Rotate a weight matrix with ConvRot and pack it as signed W4.""" + if quant_group_size != _INT4_GROUP_SIZE: + raise ValueError(f"int4 MMA kernel requires quant_group_size {_INT4_GROUP_SIZE}") + validate_w4a4_shape(weight, convrot_groupsize, quant_group_size) + h = _build_hadamard(convrot_groupsize, device=weight.device, dtype=weight.dtype) + weight_rot = _rotate_weight(weight, h, convrot_groupsize) + return quantize_signed_int4_rowwise(weight_rot, stochastic_rounding=stochastic_rounding) + + +def dequantize_convrot_w4a4_weight( + qdata: torch.Tensor, + scales: torch.Tensor, + convrot_groupsize: int = 256, + quant_group_size: int = _INT4_GROUP_SIZE, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Dequantize packed ConvRot W4A4 weights and rotate back to original basis.""" + if quant_group_size != _INT4_GROUP_SIZE: + raise ValueError(f"int4 MMA kernel requires quant_group_size {_INT4_GROUP_SIZE}") + w_int = _unpack_int4_row_major(qdata).to(torch.float32) + w_rot = w_int * scales.to(device=qdata.device, dtype=torch.float32).reshape(-1, 1) + h = _build_hadamard(convrot_groupsize, device=qdata.device, dtype=torch.float32) + return _rotate_weight(w_rot.float(), h, convrot_groupsize).to(output_dtype) + + +def int4_linear( + x_qdata: torch.Tensor, + weight: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + x_int = _unpack_int4_row_major(x_qdata).to(dtype=out_dtype) + w_int = _unpack_int4_row_major(weight).to(dtype=out_dtype) + out = x_int @ w_int.t() + out = out * x_scale.reshape(-1, 1).to(device=out.device, dtype=out_dtype) + out = out * weight_scale.reshape(1, -1).to(device=out.device, dtype=out_dtype) + if bias is not None: + out = out + bias.to(device=out.device, dtype=out_dtype).reshape(1, -1) + return out + + +def convrot_w4a4_linear( + x: torch.Tensor, + qweight: torch.Tensor, + wscales: torch.Tensor, + bias: torch.Tensor | None = None, + convrot_groupsize: int = 256, + quant_group_size: int = _INT4_GROUP_SIZE, +) -> torch.Tensor: + """Compute ``x @ W.T + bias`` using the eager ConvRot W4A4 path.""" + if quant_group_size != _INT4_GROUP_SIZE: + raise ValueError(f"int4 MMA kernel requires quant_group_size {_INT4_GROUP_SIZE}") + if x.shape[-1] != qweight.shape[-1] * 2: + raise ValueError(f"Input K={x.shape[-1]} does not match qweight K={qweight.shape[-1] * 2}") + if x.shape[-1] % convrot_groupsize != 0: + raise ValueError(f"Input K={x.shape[-1]} not divisible by convrot_groupsize {convrot_groupsize}") + + orig_shape = x.shape + x2d = x.reshape(-1, orig_shape[-1]).contiguous() + h = _build_hadamard(convrot_groupsize, device=x2d.device, dtype=x2d.dtype) + x_rot = _rotate_activation(x2d, h, convrot_groupsize).contiguous() + qact, x_scale = quantize_signed_int4_rowwise(x_rot) + out = int4_linear(qact, qweight, x_scale, wscales, bias, x.dtype) + return out[: x2d.shape[0]].reshape(*orig_shape[:-1], qweight.shape[0]) diff --git a/comfy_kitchen/tensor/__init__.py b/comfy_kitchen/tensor/__init__.py index e1a9bf5..aa0efd8 100644 --- a/comfy_kitchen/tensor/__init__.py +++ b/comfy_kitchen/tensor/__init__.py @@ -10,6 +10,12 @@ register_layout_class, register_layout_op, ) +from .convrot_w4a4 import ( + TensorCoreConvRotW4A4Layout, + convrot_w4a4_linear, + dequantize_convrot_w4a4_weight, + quantize_convrot_w4a4_weight, +) from .fp8 import TensorCoreFP8Layout from .int8 import TensorWiseINT8Layout from .mxfp8 import TensorCoreMXFP8Layout @@ -27,16 +33,20 @@ "QuantizedLayout", "QuantizedTensor", "TensorCoreAWQW4A16Layout", + "TensorCoreConvRotW4A4Layout", "TensorCoreFP8Layout", "TensorCoreMXFP8Layout", "TensorCoreNVFP4Layout", "TensorWiseINT8Layout", "TensorCoreSVDQuantW4A4Layout", + "convrot_w4a4_linear", "dequantize_args", + "dequantize_convrot_w4a4_weight", "get_cuda_capability", "get_layout_class", "register_layout_class", "register_layout_op", + "quantize_convrot_w4a4_weight", "svdquant_w4a4_can_share_quant", "svdquant_w4a4_fuse_linear_weights", "svdquant_w4a4_fused_grouped_linear", @@ -44,6 +54,7 @@ ] register_layout_class("TensorCoreAWQW4A16Layout", TensorCoreAWQW4A16Layout) +register_layout_class("TensorCoreConvRotW4A4Layout", TensorCoreConvRotW4A4Layout) register_layout_class("TensorCoreFP8Layout", TensorCoreFP8Layout) register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout) register_layout_class("TensorCoreNVFP4Layout", TensorCoreNVFP4Layout) diff --git a/comfy_kitchen/tensor/convrot_w4a4.py b/comfy_kitchen/tensor/convrot_w4a4.py new file mode 100644 index 0000000..91a532c --- /dev/null +++ b/comfy_kitchen/tensor/convrot_w4a4.py @@ -0,0 +1,244 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 Comfy Org. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""ConvRot W4A4 quantization layout. + +This is the plain ConvRot path from the paper adapted to kitchen's existing +int4 tensor-core GEMM: offline regular-Hadamard weight rotation, online +regular-Hadamard activation rotation, symmetric int4 quantization, and int4 +MMA with per-group scales. +""" +from __future__ import annotations + +import dataclasses +from dataclasses import dataclass + +import torch + +from comfy_kitchen.backends.eager.svdquant import _INT4_GROUP_SIZE +from comfy_kitchen.registry import registry + +from .base import ( + BaseLayoutParams, + QuantizedLayout, + QuantizedTensor, + dequantize_args, + register_layout_op, +) + + +def quantize_convrot_w4a4_weight( + weight: torch.Tensor, + convrot_groupsize: int = 256, + quant_group_size: int = _INT4_GROUP_SIZE, + stochastic_rounding: int | None = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Rotate a weight matrix with ConvRot and pack it as signed W4.""" + impl = registry.get_implementation( + "quantize_convrot_w4a4_weight", + kwargs={ + "weight": weight, + "convrot_groupsize": convrot_groupsize, + "quant_group_size": quant_group_size, + "stochastic_rounding": stochastic_rounding, + }, + ) + return impl( + weight, + convrot_groupsize=convrot_groupsize, + quant_group_size=quant_group_size, + stochastic_rounding=stochastic_rounding, + ) + + +def dequantize_convrot_w4a4_weight( + qdata: torch.Tensor, + scales: torch.Tensor, + convrot_groupsize: int = 256, + quant_group_size: int = _INT4_GROUP_SIZE, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Dequantize packed ConvRot W4A4 weights and rotate back to original basis.""" + impl = registry.get_implementation( + "dequantize_convrot_w4a4_weight", + kwargs={ + "qdata": qdata, + "scales": scales, + "convrot_groupsize": convrot_groupsize, + "quant_group_size": quant_group_size, + "output_dtype": output_dtype, + }, + ) + return impl( + qdata, + scales, + convrot_groupsize=convrot_groupsize, + quant_group_size=quant_group_size, + output_dtype=output_dtype, + ) + + +def convrot_w4a4_linear( + x: torch.Tensor, + qweight: torch.Tensor, + wscales: torch.Tensor, + bias: torch.Tensor | None = None, + convrot_groupsize: int = 256, + quant_group_size: int = _INT4_GROUP_SIZE, +) -> torch.Tensor: + """Compute ``x @ W.T + bias`` using ConvRot W4A4 int4 MMA.""" + impl = registry.get_implementation( + "convrot_w4a4_linear", + kwargs={ + "x": x, + "qweight": qweight, + "wscales": wscales, + "bias": bias, + "convrot_groupsize": convrot_groupsize, + "quant_group_size": quant_group_size, + }, + ) + return impl( + x, + qweight, + wscales, + bias=bias, + convrot_groupsize=convrot_groupsize, + quant_group_size=quant_group_size, + ) + + +class TensorCoreConvRotW4A4Layout(QuantizedLayout): + """ConvRot W4A4 weight layout using kitchen's int4 tensor-core GEMM.""" + + MIN_SM_VERSION = (8, 0) + QUANTIZES_INPUT = False + + @dataclass(frozen=True) + class Params(BaseLayoutParams): + convrot_groupsize: int = 256 + quant_group_size: int = _INT4_GROUP_SIZE + transposed: bool = False + + def _tensor_fields(self) -> list[str]: + return ["scale"] + + def _validate_tensor_fields(self): + return + + @classmethod + def quantize( + cls, + tensor: torch.Tensor, + convrot_groupsize: int = 256, + quant_group_size: int = _INT4_GROUP_SIZE, + stochastic_rounding: int | None = 0, + **kwargs, + ) -> tuple[torch.Tensor, Params]: + qdata, scales = quantize_convrot_w4a4_weight( + tensor, + convrot_groupsize, + quant_group_size, + stochastic_rounding=stochastic_rounding, + ) + params = cls.Params( + scale=scales, + orig_dtype=tensor.dtype, + orig_shape=tuple(tensor.shape), + convrot_groupsize=convrot_groupsize, + quant_group_size=quant_group_size, + ) + return qdata, params + + @classmethod + def dequantize(cls, qdata: torch.Tensor, params: Params) -> torch.Tensor: + return dequantize_convrot_w4a4_weight( + qdata, + params.scale, + params.convrot_groupsize, + params.quant_group_size, + params.orig_dtype, + ) + + @classmethod + def get_plain_tensors(cls, qtensor: QuantizedTensor) -> tuple[torch.Tensor, torch.Tensor]: + return qtensor._qdata, qtensor._params.scale + + @classmethod + def state_dict_tensors(cls, qdata: torch.Tensor, params: Params) -> dict[str, torch.Tensor]: + return {"": qdata, "_scale": params.scale} + + @classmethod + def requantize_kwargs(cls, qtensor: QuantizedTensor) -> dict[str, object]: + params = qtensor._params + return { + "convrot_groupsize": params.convrot_groupsize, + "quant_group_size": params.quant_group_size, + } + + +@register_layout_op(torch.ops.aten.t.default, TensorCoreConvRotW4A4Layout) +def _handle_convrot_w4a4_t(qt, args, kwargs): + input_tensor = args[0] + if not isinstance(input_tensor, QuantizedTensor): + return torch.ops.aten.t.default(*args, **kwargs) + old = input_tensor._params + new_params = dataclasses.replace( + old, + orig_shape=(old.orig_shape[1], old.orig_shape[0]), + transposed=not old.transposed, + ) + return QuantizedTensor(input_tensor._qdata, "TensorCoreConvRotW4A4Layout", new_params) + + +def _resolve_convrot_w4a4_rhs(rhs: QuantizedTensor) -> QuantizedTensor: + if not rhs._params.transposed: + raise RuntimeError("ConvRot W4A4 GEMM expects RHS W.T. Use F.linear(x, W) or mm(x, W.t()).") + return rhs + + +def _convrot_w4a4_forward(input_tensor: torch.Tensor, weight: QuantizedTensor, bias: torch.Tensor | None): + qweight, wscales = TensorCoreConvRotW4A4Layout.get_plain_tensors(weight) + params = weight._params + return convrot_w4a4_linear( + input_tensor, + qweight, + wscales, + bias=bias, + convrot_groupsize=params.convrot_groupsize, + quant_group_size=params.quant_group_size, + ) + + +@register_layout_op(torch.ops.aten.linear.default, TensorCoreConvRotW4A4Layout) +def _handle_convrot_w4a4_linear(qt, args, kwargs): + input_tensor, weight = args[0], args[1] + bias = args[2] if len(args) > 2 else None + if not isinstance(weight, QuantizedTensor): + return torch.nn.functional.linear(*dequantize_args((input_tensor, weight, bias))) + if isinstance(input_tensor, QuantizedTensor): + input_tensor = input_tensor.dequantize() + if weight._params.transposed: + return torch.nn.functional.linear(input_tensor, weight.dequantize(), bias) + return _convrot_w4a4_forward(input_tensor, weight, bias) + + +@register_layout_op(torch.ops.aten.mm.default, TensorCoreConvRotW4A4Layout) +def _handle_convrot_w4a4_mm(qt, args, kwargs): + a, b = args[0], args[1] + if not isinstance(b, QuantizedTensor): + return torch.mm(*dequantize_args((a, b))) + if isinstance(a, QuantizedTensor): + a = a.dequantize() + b = _resolve_convrot_w4a4_rhs(b) + return _convrot_w4a4_forward(a, b, bias=None) + + +@register_layout_op(torch.ops.aten.addmm.default, TensorCoreConvRotW4A4Layout) +def _handle_convrot_w4a4_addmm(qt, args, kwargs): + bias, a, b = args[0], args[1], args[2] + if not isinstance(b, QuantizedTensor): + return torch.addmm(*dequantize_args((bias, a, b))) + if isinstance(a, QuantizedTensor): + a = a.dequantize() + b = _resolve_convrot_w4a4_rhs(b) + return _convrot_w4a4_forward(a, b, bias=bias) From cc511a2393132ceaf9cee9e7bc470c9bdbda59de Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 8 Jul 2026 17:55:53 -0400 Subject: [PATCH 2/7] Overflow fix on int4. --- .../backends/cuda/ops/convrot_w4a4.cu | 28 ++++++++++++++++--- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu b/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu index 94703e2..5eea776 100644 --- a/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu +++ b/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu @@ -7,6 +7,7 @@ #include #include #include +#include #include #include @@ -58,6 +59,17 @@ __device__ __forceinline__ int quantize_int4_value(float x, float inv_scale, uin return max(-kInt4Max, min(kInt4Max, q)); } +template +__device__ __forceinline__ float finite_max_for_dtype(); +template<> __device__ __forceinline__ float finite_max_for_dtype() { return FLT_MAX; } +template<> __device__ __forceinline__ float finite_max_for_dtype<__half>() { return 65504.0f; } +template<> __device__ __forceinline__ float finite_max_for_dtype<__nv_bfloat16>() { return 3.38953139e38f; } + +template +__device__ __forceinline__ float finite_absmax_for_int4_scale(float abs_max) { + return fminf(abs_max, finite_max_for_dtype()); +} + __device__ __forceinline__ int unpack_int4_nibble(uint32_t v) { return static_cast((v & 0x0fu) ^ 0x08u) - 8; } @@ -1000,7 +1012,9 @@ __global__ void quantize_int4_rowwise_kernel( absmax = fmaxf(absmax, fabsf(to_float(x[row_offset + col]))); } absmax = block_reduce_max(absmax, warp_smem, &block_smem); - const float scale = fmaxf(absmax * (1.0f / static_cast(kInt4Max)), 1.0e-10f); + const float scale = fmaxf( + finite_absmax_for_int4_scale(absmax) * (1.0f / static_cast(kInt4Max)), + 1.0e-10f); if (tid == 0) { scales[row] = scale; } @@ -1106,7 +1120,9 @@ __global__ void quantize_int4_rowwise_convrot64_kernel( } abs_max = block_reduce_max(abs_max, warp_smem, &block_smem); - const float scale = fmaxf(abs_max * (1.0f / static_cast(kInt4Max)), 1.0e-10f); + const float scale = fmaxf( + finite_absmax_for_int4_scale(abs_max) * (1.0f / static_cast(kInt4Max)), + 1.0e-10f); if (tid == 0) { scales[row] = scale; } @@ -1201,7 +1217,9 @@ __global__ void quantize_int4_rowwise_convrot64_to_int8_float_kernel( } abs_max = block_reduce_max(abs_max, warp_smem, &block_smem); - const float scale = fmaxf(abs_max * (1.0f / static_cast(kInt4Max)), 1.0e-10f); + const float scale = fmaxf( + finite_absmax_for_int4_scale(abs_max) * (1.0f / static_cast(kInt4Max)), + 1.0e-10f); if (tid == 0) { scales[row] = scale; } @@ -1287,7 +1305,9 @@ __global__ void quantize_int4_rowwise_convrot_small_kernel( } abs_max = block_reduce_max(abs_max, warp_smem, &block_smem); - const float scale = fmaxf(abs_max * (1.0f / static_cast(kInt4Max)), 1.0e-10f); + const float scale = fmaxf( + finite_absmax_for_int4_scale(abs_max) * (1.0f / static_cast(kInt4Max)), + 1.0e-10f); if (tid == 0) { scales[row] = scale; } From 4e2518f73919a594882b4e91a8d13913a7a30839 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 8 Jul 2026 18:56:45 -0400 Subject: [PATCH 3/7] Missing tests. --- tests/test_convrot_w4a4.py | 171 +++++++++++++++++++++++++++++++++++++ 1 file changed, 171 insertions(+) create mode 100644 tests/test_convrot_w4a4.py diff --git a/tests/test_convrot_w4a4.py b/tests/test_convrot_w4a4.py new file mode 100644 index 0000000..dd2345e --- /dev/null +++ b/tests/test_convrot_w4a4.py @@ -0,0 +1,171 @@ +"""Tests for ConvRot W4A4 int4 tensor-core layout.""" +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as functional + +import comfy_kitchen as ck +from comfy_kitchen.backends import cuda as cuda_backend +from comfy_kitchen.backends.eager.svdquant import _unpack_int4_row_major +from comfy_kitchen.tensor import QuantizedTensor +from comfy_kitchen.tensor.convrot_w4a4 import ( + convrot_w4a4_linear, + dequantize_convrot_w4a4_weight, + quantize_convrot_w4a4_weight, +) + + +def test_convrot_w4a4_weight_quantize_contract(seed): + w = torch.randn(16, 256, dtype=torch.float32) + + q, scale = quantize_convrot_w4a4_weight(w) + q_values = _unpack_int4_row_major(q) + dq = dequantize_convrot_w4a4_weight(q, scale, output_dtype=w.dtype) + + assert q.shape == (16, 128) + assert scale.shape == (16,) + assert q_values.min().item() >= -7 + assert q_values.max().item() <= 7 + assert dq.shape == w.shape + + rel_err = (w - dq).abs() / (w.abs().max() + 1e-8) + assert rel_err.mean().item() < 0.04 + + +def test_convrot_w4a4_stochastic_rounding_eager(seed): + torch.manual_seed(1234) + w = torch.randn(16, 256, dtype=torch.float32) + + with ck.registry.use_backend("eager"): + q1, scale1 = quantize_convrot_w4a4_weight(w, stochastic_rounding=123) + q2, scale2 = quantize_convrot_w4a4_weight(w, stochastic_rounding=123) + q3, scale3 = quantize_convrot_w4a4_weight(w, stochastic_rounding=124) + + assert torch.equal(q1, q2) + assert not torch.equal(q1, q3) + assert torch.equal(scale1, scale2) + assert torch.equal(scale1, scale3) + + +def test_convrot_w4a4_linear_eager_matches_dequantized_weight(seed): + x = torch.randn(5, 256, dtype=torch.float32) + w = torch.randn(17, 256, dtype=torch.float32) + bias = torch.randn(17, dtype=torch.float32) + q, scale = quantize_convrot_w4a4_weight(w) + + with ck.registry.use_backend("eager"): + out = convrot_w4a4_linear(x, q, scale, bias=bias) + + ref = functional.linear(x, w, bias) + rel_err = (out - ref).abs() / (ref.abs().max() + 1e-8) + assert rel_err.mean().item() < 0.08 + + +def test_convrot_w4a4_layout_linear_mm_addmm(seed): + x = torch.randn(5, 256, dtype=torch.float32) + w = torch.randn(17, 256, dtype=torch.float32) + bias = torch.randn(17, dtype=torch.float32) + qt = QuantizedTensor.from_float(w, "TensorCoreConvRotW4A4Layout") + + out_linear = functional.linear(x, qt, bias) + out_mm = torch.mm(x, qt.t()) + out_addmm = torch.addmm(bias, x, qt.t()) + + assert out_linear.shape == (5, 17) + assert out_mm.shape == (5, 17) + assert out_addmm.shape == (5, 17) + assert torch.allclose(out_linear, out_mm + bias, rtol=1e-5, atol=1e-5) + assert torch.allclose(out_linear, out_addmm, rtol=1e-5, atol=1e-5) + + +def test_convrot_w4a4_rejects_bad_groups(seed): + w = torch.randn(16, 250, dtype=torch.float32) + with pytest.raises(ValueError, match="not divisible by convrot_groupsize"): + quantize_convrot_w4a4_weight(w) + + +def test_convrot_w4a4_cuda_smoke(seed): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + x = torch.randn(64, 256, device="cuda", dtype=torch.bfloat16) + w = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + bias = torch.randn(128, device="cuda", dtype=torch.bfloat16) + qt = QuantizedTensor.from_float(w, "TensorCoreConvRotW4A4Layout") + + dq = qt.dequantize() + out = functional.linear(x, qt, bias) + assert dq.shape == w.shape + assert dq.device.type == "cuda" + assert dq.dtype == torch.bfloat16 + assert out.shape == (64, 128) + assert out.dtype == torch.bfloat16 + + +def test_convrot_w4a4_cuda_no_bias_large_m(seed): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + x = torch.randn(1152, 256, device="cuda", dtype=torch.bfloat16) + w = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16) + qt = QuantizedTensor.from_float(w, "TensorCoreConvRotW4A4Layout") + + out = functional.linear(x, qt) + assert out.shape == (1152, 128) + assert out.dtype == torch.bfloat16 + + +def test_convrot_w4a4_cuda_large_k_quantize_matches_reference(seed): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + x = torch.randn(2, 15360, device="cuda", dtype=torch.bfloat16) + q_cuda, scale_cuda = cuda_backend.quantize_int4_rowwise_convrot64(x, 256) + q_values = _unpack_int4_row_major(q_cuda) + dq = cuda_backend.dequantize_convrot_w4a4_weight(q_cuda, scale_cuda.reshape(-1), output_dtype=torch.float32) + + assert q_cuda.shape == (2, 7680) + assert scale_cuda.shape == (2, 1) + assert q_values.min().item() >= -7 + assert q_values.max().item() <= 7 + assert dq.shape == x.shape + rel_err = (x.float() - dq).abs() / (x.float().abs().max() + 1e-8) + assert rel_err.mean().item() < 0.04 + + +@pytest.mark.parametrize("convrot_groupsize", [16, 64, 256]) +def test_convrot_w4a4_cuda_quantize_clamps_overflow_scale(convrot_groupsize): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + x = torch.empty(2, 256, device="cuda", dtype=torch.float16) + pattern = torch.tensor([65504.0, 65504.0, 65504.0, -65504.0], device="cuda", dtype=torch.float16) + x.copy_(pattern.repeat(x.numel() // pattern.numel()).reshape_as(x)) + + q, scale = cuda_backend.quantize_int4_rowwise_convrot64(x, convrot_groupsize) + q_values = _unpack_int4_row_major(q) + + assert q.shape == (2, 128) + assert scale.shape == (2, 1) + assert torch.isfinite(scale).all() + assert q_values.min().item() >= -7 + assert q_values.max().item() <= 7 + + +def test_convrot_w4a4_stochastic_rounding_cuda(seed): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + torch.manual_seed(1234) + w = torch.randn(16, 256, device="cuda", dtype=torch.bfloat16) + + with ck.registry.use_backend("cuda"): + q1, scale1 = quantize_convrot_w4a4_weight(w, stochastic_rounding=123) + q2, scale2 = quantize_convrot_w4a4_weight(w, stochastic_rounding=123) + q3, scale3 = quantize_convrot_w4a4_weight(w, stochastic_rounding=124) + + assert torch.equal(q1, q2) + assert not torch.equal(q1, q3) + assert torch.equal(scale1, scale2) + assert torch.equal(scale1, scale3) From 769f3249a9b4e8a13121ba6b2276d3c4323a5abc Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 8 Jul 2026 22:59:33 -0400 Subject: [PATCH 4/7] bring up int4 to same speed as int8 on blackwell --- comfy_kitchen/backends/cuda/__init__.py | 140 ++++++- .../backends/cuda/dlpack_bindings.cpp | 197 +++++++++ .../backends/cuda/ops/convrot_w4a4.cu | 376 +++++++++++++++++- .../backends/cuda/ops/cutlass_gemm_int8.cu | 136 ++++++- 4 files changed, 823 insertions(+), 26 deletions(-) diff --git a/comfy_kitchen/backends/cuda/__init__.py b/comfy_kitchen/backends/cuda/__init__.py index 5780dbf..441e7ac 100644 --- a/comfy_kitchen/backends/cuda/__init__.py +++ b/comfy_kitchen/backends/cuda/__init__.py @@ -161,6 +161,8 @@ def find_lib_dir(start_dir, lib_pattern): _cutlass_int8_device_cache: dict[int, bool] = {} _shared_memory_per_block_cache: dict[int, int] = {} _FORCE_INT4_INT8_FALLBACK = os.environ.get("COMFY_KITCHEN_FORCE_INT4_INT8_FALLBACK", "0") == "1" +_INT4_PACKED_WEIGHT_SMALL_M_MAX = 8 +_INT4_INT8_WEIGHT_CHUNK_N = max(1, int(os.environ.get("COMFY_KITCHEN_INT4_INT8_WEIGHT_CHUNK_N", "4096"))) def _cuda_device_is_turing(device_index: int) -> bool: @@ -595,6 +597,103 @@ def _unpack_int4_to_int8_cuda(qdata: torch.Tensor) -> torch.Tensor: return output +def _int4_weight_int8_act_gemv_dequant( + x_int8: torch.Tensor, + weight_packed: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + if x_int8.dim() != 2: + raise ValueError("packed INT4 weight GEMV expects a 2D INT8 activation") + if weight_packed.dim() != 2 or x_int8.shape[1] != weight_packed.shape[1] * 2: + raise ValueError("packed INT4 weight GEMV K dimensions do not match") + + m = x_int8.shape[0] + n = weight_packed.shape[0] + output = torch.empty((m, n), dtype=out_dtype, device=x_int8.device) + x_scale_arg = x_scale.to(device=x_int8.device, dtype=torch.float32).reshape(-1, 1).contiguous() + if x_scale_arg.shape[0] != m: + raise ValueError(f"packed INT4 weight GEMV x_scale must have {m} values, got {x_scale_arg.shape[0]}") + weight_scale_arg = weight_scale.to(device=x_int8.device, dtype=torch.float32).reshape(-1).contiguous() + if weight_scale_arg.numel() != n: + raise ValueError(f"packed INT4 weight GEMV weight_scale must have {n} values, got {weight_scale_arg.numel()}") + bias_arg = bias if bias is not None else _empty_cuda_tensor(x_int8.device, out_dtype) + if bias is not None and (bias.device != x_int8.device or bias.dtype != out_dtype or not bias.is_contiguous()): + bias_arg = bias.to(device=x_int8.device, dtype=out_dtype).contiguous() + + stream_ptr = torch.cuda.current_stream(x_int8.device).cuda_stream + _C.int4_weight_int8_act_gemv_dequant( + _wrap_for_dlpack(x_int8), + _wrap_for_dlpack(weight_packed.contiguous()), + _wrap_for_dlpack(x_scale_arg), + _wrap_for_dlpack(weight_scale_arg), + _wrap_for_dlpack(bias_arg), + _wrap_for_dlpack(output), + DTYPE_TO_CODE[out_dtype], + stream_ptr, + ) + return output + + +def _int4_int8_weight_chunk_cols(m: int, n: int) -> int: + if n <= 2560: + return n + if m <= 128: + return min(n, _INT4_INT8_WEIGHT_CHUNK_N) + return min(n, _INT4_INT8_WEIGHT_CHUNK_N) + + +def _int4_weight_int8_act_gemm_dequant_chunked( + x_int8: torch.Tensor, + weight_packed: torch.Tensor, + x_scale: torch.Tensor, + weight_scale: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + if x_int8.dim() != 2: + raise ValueError("chunked INT4 weight GEMM expects a 2D INT8 activation") + if weight_packed.dim() != 2 or x_int8.shape[1] != weight_packed.shape[1] * 2: + raise ValueError("chunked INT4 weight GEMM K dimensions do not match") + if not x_int8.is_contiguous() or not weight_packed.is_contiguous(): + raise ValueError("chunked INT4 weight GEMM expects contiguous activation and weight tensors") + + m, k = x_int8.shape + n = weight_packed.shape[0] + chunk_cols = _int4_int8_weight_chunk_cols(m, n) + output = torch.empty((m, n), dtype=out_dtype, device=x_int8.device) + weight_workspace = torch.empty((chunk_cols, k), dtype=torch.int8, device=x_int8.device) + acc_workspace = torch.empty((m, chunk_cols), dtype=torch.int32, device=x_int8.device) + x_scale_arg = x_scale.to(device=x_int8.device, dtype=torch.float32).reshape(-1, 1).contiguous() + if x_scale_arg.shape[0] != m: + raise ValueError(f"chunked INT4 weight GEMM x_scale must have {m} values, got {x_scale_arg.shape[0]}") + weight_scale_arg = weight_scale.to(device=x_int8.device, dtype=torch.float32).reshape(-1).contiguous() + if weight_scale_arg.numel() != n: + raise ValueError(f"chunked INT4 weight GEMM weight_scale must have {n} values, got {weight_scale_arg.numel()}") + bias_arg = bias if bias is not None else _empty_cuda_tensor(x_int8.device, out_dtype) + if bias is not None: + bias_arg = bias.to(device=x_int8.device, dtype=torch.float32).contiguous() + + stream_ptr = torch.cuda.current_stream(x_int8.device).cuda_stream + _C.int4_weight_int8_act_gemm_dequant_chunked( + _wrap_for_dlpack(x_int8), + _wrap_for_dlpack(weight_packed), + _wrap_for_dlpack(x_scale_arg), + _wrap_for_dlpack(weight_scale_arg), + _wrap_for_dlpack(bias_arg), + _wrap_for_dlpack(output), + _wrap_for_dlpack(weight_workspace), + _wrap_for_dlpack(acc_workspace), + _wrap_for_dlpack(get_cublas_workspace()), + chunk_cols, + DTYPE_TO_CODE[out_dtype], + stream_ptr, + ) + return output + + def _int4_linear_via_int8_values( x_int8: torch.Tensor, weight_int8: torch.Tensor, @@ -621,13 +720,26 @@ def _int4_linear_via_int8_values( if bias is not None and (bias.device != x_int8.device or bias.dtype != out_dtype or not bias.is_contiguous()): bias_arg = bias.to(device=x_int8.device, dtype=out_dtype).contiguous() + stream_ptr = torch.cuda.current_stream(x_int8.device).cuda_stream + if m == 1 and k % 4 == 0 and hasattr(_C, "int8_gemv_dequant"): + _C.int8_gemv_dequant( + _wrap_for_dlpack(x_int8), + _wrap_for_dlpack(weight_int8), + _wrap_for_dlpack(x_scale_arg.reshape(1, 1)), + _wrap_for_dlpack(weight_scale_arg), + _wrap_for_dlpack(bias_arg), + _wrap_for_dlpack(output), + DTYPE_TO_CODE[out_dtype], + stream_ptr, + ) + return output + used_cutlass = False + prefer_cublas_fallback = _prefer_cublas_int8_fallback(m, n, k) if ( - not _DISABLE_CUTLASS_INT8 + not prefer_cublas_fallback + and not _DISABLE_CUTLASS_INT8 and _cuda_device_supports_cutlass_int8_dequant(x_int8) - and m <= 512 - and n <= k - and k <= 4096 ): ws_cutlass = weight_scale_arg if weight_scale_arg.numel() == n else weight_scale_arg.expand(n).contiguous() bias_f32 = bias_arg.to(torch.float32).contiguous() if bias is not None else bias_arg @@ -655,7 +767,6 @@ def _int4_linear_via_int8_values( cublas_weight = weight_int8 out_int32 = torch.empty((m, padded_n), dtype=torch.int32, device=x_int8.device) - stream_ptr = torch.cuda.current_stream(x_int8.device).cuda_stream _C.cublas_gemm_int8( _wrap_for_dlpack(cublas_x), _wrap_for_dlpack(cublas_weight), @@ -690,8 +801,10 @@ def _int4_linear_via_int8( n = weight.shape[0] k = k_half * 2 x_int8 = _unpack_int4_to_int8_cuda(x_qdata) + if x_int8.shape[0] <= _INT4_PACKED_WEIGHT_SMALL_M_MAX and hasattr(_C, "int4_weight_int8_act_gemv_dequant"): + return _int4_weight_int8_act_gemv_dequant(x_int8, weight, x_scale, weight_scale, bias, out_dtype) if weight_int8 is None: - weight_int8 = _unpack_int4_to_int8_cuda(weight) + return _int4_weight_int8_act_gemm_dequant_chunked(x_int8, weight, x_scale, weight_scale, bias, out_dtype) elif weight_int8.shape != (n, k) or weight_int8.dtype != torch.int8 or weight_int8.device != weight.device: raise ValueError("prepared INT8 fallback weight has incompatible shape, dtype, or device") return _int4_linear_via_int8_values(x_int8, weight_int8, x_scale, weight_scale, bias, out_dtype) @@ -852,7 +965,6 @@ def convrot_w4a4_linear( orig_shape = x.shape x2d = x.reshape(-1, orig_shape[-1]).contiguous() if not _cuda_device_supports_native_int4_mma(x2d): - weight_int8 = prepare_int4_weight_for_int8_linear(qweight) if ( convrot_groupsize == 256 and x2d.shape[-1] % 256 == 0 @@ -865,9 +977,19 @@ def convrot_w4a4_linear( else: h = _build_hadamard(convrot_groupsize, device=x2d.device, dtype=x2d.dtype) qact_int8, x_scale = quantize_and_rotate_rowwise(x2d, h, convrot_groupsize) - out = _int4_linear_via_int8_values( + if qact_int8.shape[0] <= _INT4_PACKED_WEIGHT_SMALL_M_MAX and hasattr(_C, "int4_weight_int8_act_gemv_dequant"): + out = _int4_weight_int8_act_gemv_dequant( + qact_int8, + qweight, + x_scale, + wscales, + bias, + x.dtype, + ) + return out.reshape(*orig_shape[:-1], qweight.shape[0]) + out = _int4_weight_int8_act_gemm_dequant_chunked( qact_int8, - weight_int8, + qweight, x_scale, wscales, bias, diff --git a/comfy_kitchen/backends/cuda/dlpack_bindings.cpp b/comfy_kitchen/backends/cuda/dlpack_bindings.cpp index dc538a5..ce65aa9 100644 --- a/comfy_kitchen/backends/cuda/dlpack_bindings.cpp +++ b/comfy_kitchen/backends/cuda/dlpack_bindings.cpp @@ -818,6 +818,43 @@ extern "C" { int64_t K_half, cudaStream_t stream); + void launch_int4_weight_int8_act_gemv_dequant_kernel( + const void* input, + const void* weight, + const void* x_scales, + const void* weight_scales, + const void* bias, + void* output, + int64_t num_rows, + int64_t num_cols, + int64_t K, + int64_t weight_scale_size, + bool has_bias, + int output_dtype_code, + int bias_dtype_code, + cudaStream_t stream); + + void launch_int4_weight_int8_act_gemm_dequant_chunked_kernel( + const void* input, + const void* weight, + const void* x_scales, + const void* weight_scales, + const void* bias, + void* output, + void* weight_workspace, + void* acc_workspace, + void* cublas_workspace, + int64_t cublas_workspace_size, + int64_t num_rows, + int64_t num_cols, + int64_t K, + int64_t weight_scale_size, + int64_t chunk_cols, + bool has_bias, + int output_dtype_code, + int bias_dtype_code, + cudaStream_t stream); + bool launch_cutlass_int8_dequant( const void* A, const void* B, @@ -1256,6 +1293,140 @@ void unpack_int4_to_int8( launch_unpack_int4_to_int8_kernel(input.data(), output.data(), rows, K_half, stream); } +void int4_weight_int8_act_gemv_dequant( + nb::ndarray, nb::device::cuda> input, + nb::ndarray, nb::device::cuda> weight, + nb::ndarray, nb::device::cuda> x_scales, + nb::ndarray weight_scales, + nb::ndarray bias, + nb::ndarray, nb::device::cuda> output, + int output_dtype_code, + uintptr_t stream_ptr) { + + const int64_t M = input.shape(0); + const int64_t K = input.shape(1); + const int64_t N = weight.shape(0); + if (weight.shape(1) * 2 != K) { + throw std::runtime_error("packed INT4 weight GEMV weight K mismatch"); + } + if (x_scales.shape(0) != M || x_scales.shape(1) != 1) { + throw std::runtime_error("packed INT4 weight GEMV activation scale shape mismatch"); + } + if (output.shape(0) != M || output.shape(1) != N) { + throw std::runtime_error("packed INT4 weight GEMV output shape mismatch"); + } + if (output_dtype_code < 0 || output_dtype_code > 2) { + throw std::runtime_error("Invalid packed INT4 weight GEMV output dtype code"); + } + + const bool has_bias = bias.data() && bias.size() > 0; + int bias_dtype_code = output_dtype_code; + if (has_bias) { + if (bias.shape(0) != N) { + throw std::runtime_error("packed INT4 weight GEMV bias shape mismatch"); + } + bias_dtype_code = map_dtype_to_code(bias.dtype()); + if (bias_dtype_code < 0 || bias_dtype_code > 2) { + throw std::runtime_error("Unsupported bias dtype for packed INT4 weight GEMV"); + } + } + + cudaStream_t stream = reinterpret_cast(stream_ptr); + launch_int4_weight_int8_act_gemv_dequant_kernel( + input.data(), + weight.data(), + x_scales.data(), + weight_scales.data(), + has_bias ? bias.data() : nullptr, + output.data(), + M, + N, + K, + static_cast(weight_scales.size()), + has_bias, + output_dtype_code, + bias_dtype_code, + stream); +} + +void int4_weight_int8_act_gemm_dequant_chunked( + nb::ndarray, nb::device::cuda> input, + nb::ndarray, nb::device::cuda> weight, + nb::ndarray, nb::device::cuda> x_scales, + nb::ndarray weight_scales, + nb::ndarray bias, + nb::ndarray, nb::device::cuda> output, + nb::ndarray, nb::device::cuda> weight_workspace, + nb::ndarray, nb::device::cuda> acc_workspace, + nb::ndarray cublas_workspace, + int64_t chunk_cols, + int output_dtype_code, + uintptr_t stream_ptr) { + + const int64_t M = input.shape(0); + const int64_t K = input.shape(1); + const int64_t N = weight.shape(0); + const int64_t K_half = weight.shape(1); + if (K_half * 2 != K) { + throw std::runtime_error("chunked INT4 weight GEMM weight K mismatch"); + } + if (x_scales.shape(0) != M || x_scales.shape(1) != 1) { + throw std::runtime_error("chunked INT4 weight GEMM activation scale shape mismatch"); + } + if (output.shape(0) != M || output.shape(1) != N) { + throw std::runtime_error("chunked INT4 weight GEMM output shape mismatch"); + } + if (chunk_cols <= 0 || chunk_cols > N) { + throw std::runtime_error("chunked INT4 weight GEMM invalid chunk_cols"); + } + if (weight_workspace.shape(0) < chunk_cols || weight_workspace.shape(1) != K) { + throw std::runtime_error("chunked INT4 weight GEMM weight workspace shape mismatch"); + } + if (acc_workspace.shape(0) != M || acc_workspace.shape(1) < chunk_cols) { + throw std::runtime_error("chunked INT4 weight GEMM accumulator workspace shape mismatch"); + } + if (weight_scales.size() != 1 && static_cast(weight_scales.size()) != N) { + throw std::runtime_error("chunked INT4 weight GEMM weight scale shape mismatch"); + } + if (output_dtype_code < 0 || output_dtype_code > 2) { + throw std::runtime_error("Invalid chunked INT4 weight GEMM output dtype code"); + } + + const bool has_bias = bias.data() && bias.size() > 0; + int bias_dtype_code = output_dtype_code; + if (has_bias) { + if (bias.shape(0) != N) { + throw std::runtime_error("chunked INT4 weight GEMM bias shape mismatch"); + } + bias_dtype_code = map_dtype_to_code(bias.dtype()); + if (bias_dtype_code < 0 || bias_dtype_code > 2) { + throw std::runtime_error("Unsupported bias dtype for chunked INT4 weight GEMM"); + } + } + + cudaStream_t stream = reinterpret_cast(stream_ptr); + launch_int4_weight_int8_act_gemm_dequant_chunked_kernel( + input.data(), + weight.data(), + x_scales.data(), + weight_scales.data(), + has_bias ? bias.data() : nullptr, + output.data(), + weight_workspace.data(), + acc_workspace.data(), + cublas_workspace.data(), + static_cast(cublas_workspace.size()), + M, + N, + K, + static_cast(weight_scales.size()), + chunk_cols, + has_bias, + output_dtype_code, + bias_dtype_code, + stream); +} + // INT8 GEMM + fused dequant (D = acc * xs[m] * ws[n] + bias[n]) via CUTLASS. // Returns true on success; false means caller falls back to cuBLAS + dequant. bool cutlass_int8_dequant( @@ -1847,6 +2018,32 @@ NB_MODULE(_C, m) { nb::arg("output"), nb::arg("stream_ptr")); + m.def("int4_weight_int8_act_gemv_dequant", &int4_weight_int8_act_gemv_dequant, + "M=1 GEMV using INT8 activation and packed row-major INT4 weight with fused dequant", + nb::arg("input"), + nb::arg("weight"), + nb::arg("x_scales"), + nb::arg("weight_scales"), + nb::arg("bias"), + nb::arg("output"), + nb::arg("output_dtype_code"), + nb::arg("stream_ptr")); + + m.def("int4_weight_int8_act_gemm_dequant_chunked", &int4_weight_int8_act_gemm_dequant_chunked, + "Chunked INT8 GEMM using INT8 activation and packed row-major INT4 weight with fused dequant", + nb::arg("input"), + nb::arg("weight"), + nb::arg("x_scales"), + nb::arg("weight_scales"), + nb::arg("bias"), + nb::arg("output"), + nb::arg("weight_workspace"), + nb::arg("acc_workspace"), + nb::arg("cublas_workspace"), + nb::arg("chunk_cols"), + nb::arg("output_dtype_code"), + nb::arg("stream_ptr")); + m.def("cutlass_int8_dequant", &cutlass_int8_dequant, "INT8 GEMM + fused rowwise x colwise dequant + bias via CUTLASS; false -> fall back to cuBLAS", nb::arg("a"), diff --git a/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu b/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu index 5eea776..cd24b73 100644 --- a/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu +++ b/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu @@ -2,13 +2,18 @@ // // Plain ConvRot W4A4 helpers: rowwise signed int4 quantization and int4 MMA // with row/column dequant scales. +#include "dtype_dispatch.cuh" #include "svdquant_utils.cuh" #include #include #include +#include #include #include +#include +#include +#include #include #ifdef COMFY_HAVE_CUTLASS @@ -22,6 +27,31 @@ #include #endif +extern "C" void launch_cublas_gemm_int8_kernel( + const void* A_ptr, + const void* B_ptr, + void* C_ptr, + int64_t M, + int64_t N, + int64_t K, + void* workspace_ptr, + int64_t workspace_size, + cudaStream_t stream); + +extern "C" bool launch_cutlass_int8_dequant_strided( + const void* A, + const void* B, + const void* xs, + const void* ws, + const void* bias, + void* D, + int64_t M, + int64_t N, + int64_t K, + int64_t output_stride, + int out_dtype_code, + cudaStream_t stream); + namespace { using comfy::svdquant::cp_async_16b; @@ -70,6 +100,15 @@ __device__ __forceinline__ float finite_absmax_for_int4_scale(float abs_max) { return fminf(abs_max, finite_max_for_dtype()); } +inline int fp_dtype_size_bytes(int dtype_code) { + switch (dtype_code) { + case 0: return 4; + case 1: + case 2: return 2; + default: throw std::runtime_error("unsupported floating point dtype code"); + } +} + __device__ __forceinline__ int unpack_int4_nibble(uint32_t v) { return static_cast((v & 0x0fu) ^ 0x08u) - 8; } @@ -1702,22 +1741,124 @@ __global__ void int4_linear_kernel( } } -__global__ void unpack_int4_to_int8_kernel( +__device__ __forceinline__ int pack_int4_weight4_as_int8_word(uint32_t packed01, uint32_t packed23) { + const uint32_t b0 = static_cast(static_cast(unpack_int4_nibble(packed01))); + const uint32_t b1 = static_cast(static_cast(unpack_int4_nibble(packed01 >> 4))); + const uint32_t b2 = static_cast(static_cast(unpack_int4_nibble(packed23))); + const uint32_t b3 = static_cast(static_cast(unpack_int4_nibble(packed23 >> 4))); + return static_cast(b0 | (b1 << 8) | (b2 << 16) | (b3 << 24)); +} + +template +__global__ void int4_weight_int8_act_gemv_dequant_warp_kernel( + const int8_t* __restrict__ x, + const int8_t* __restrict__ weight, + const float* __restrict__ x_scales, + const float* __restrict__ weight_scales, + const BiasType* __restrict__ bias, + OutputType* __restrict__ output, + int M, + int N, + int K, + int weight_scale_size, + bool has_bias) +{ + const int lane = threadIdx.x & (kThreadsPerWarp - 1); + const int warp = threadIdx.x >> 5; + const int n = static_cast(blockIdx.x) * WARPS_PER_BLOCK + warp; + const int m = static_cast(blockIdx.y); + if (n >= N) { + return; + } + + const int K4 = K >> 2; + const int K_half = K >> 1; + const int* __restrict__ x4 = reinterpret_cast(x + static_cast(m) * K); + const int8_t* __restrict__ w_row = weight + static_cast(n) * K_half; + + int acc = 0; + for (int k4 = lane; k4 < K4; k4 += kThreadsPerWarp) { + const uint32_t packed01 = static_cast(w_row[k4 * 2]); + const uint32_t packed23 = static_cast(w_row[k4 * 2 + 1]); + acc = __dp4a(x4[k4], pack_int4_weight4_as_int8_word(packed01, packed23), acc); + } + + #pragma unroll + for (int offset = kThreadsPerWarp / 2; offset > 0; offset >>= 1) { + acc += __shfl_down_sync(0xffffffffu, acc, offset); + } + + if (lane == 0) { + const float weight_scale = weight_scales[weight_scale_size == 1 ? 0 : n]; + float value = static_cast(acc) * x_scales[m] * weight_scale; + if (has_bias) { + value += to_float(bias[n]); + } + output[static_cast(m) * N + n] = from_float(value); + } +} + +template +__global__ void dequantize_int4_weight_int8_act_chunk_kernel( + const int32_t* __restrict__ input, + const float* __restrict__ x_scales, + const float* __restrict__ weight_scales, + const BiasType* __restrict__ bias, + OutputType* __restrict__ output, + int64_t total, + int output_stride, + int chunk_cols, + int col_offset, + int weight_scale_size, + bool has_bias) +{ + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= total) { + return; + } + + const int chunk_col = static_cast(idx % chunk_cols); + const int row = static_cast(idx / chunk_cols); + const int col = col_offset + chunk_col; + const float weight_scale = weight_scales[weight_scale_size == 1 ? 0 : col]; + float value = static_cast(input[idx]) * x_scales[row] * weight_scale; + if (has_bias) { + value += to_float(bias[col]); + } + output[static_cast(row) * output_stride + col] = from_float(value); +} + +__device__ __forceinline__ uint64_t unpack_int4_u32_to_int8_u64(uint32_t packed) { + uint64_t out = 0; + #pragma unroll + for (int i = 0; i < 4; ++i) { + const uint32_t byte = (packed >> (i * 8)) & 0xffu; + const uint32_t low = static_cast(static_cast(unpack_int4_nibble(byte))); + const uint32_t high = static_cast(static_cast(unpack_int4_nibble(byte >> 4))); + out |= static_cast(low | (high << 8)) << (i * 16); + } + return out; +} + +__global__ void unpack_int4_to_int8_vec8_kernel( const int8_t* __restrict__ input, int8_t* __restrict__ output, - int64_t total_packed, - int K_half) + int64_t total_packed) { - const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (idx >= total_packed) { + const int64_t base = (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) * 8; + if (base + 7 < total_packed) { + const uint2 packed = *reinterpret_cast(input + base); + uint64_t* __restrict__ out64 = reinterpret_cast(output + base * 2); + out64[0] = unpack_int4_u32_to_int8_u64(packed.x); + out64[1] = unpack_int4_u32_to_int8_u64(packed.y); return; } - const int row = static_cast(idx / K_half); - const int packed_col = static_cast(idx - static_cast(row) * K_half); - const uint32_t packed = static_cast(input[idx]); - int8_t* out_row = output + static_cast(row) * (K_half * 2); - out_row[packed_col * 2] = static_cast(unpack_int4_nibble(packed)); - out_row[packed_col * 2 + 1] = static_cast(unpack_int4_nibble(packed >> 4)); + + for (int64_t idx = base; idx < total_packed; ++idx) { + const uint32_t packed = static_cast(input[idx]); + output[idx * 2] = static_cast(unpack_int4_nibble(packed)); + output[idx * 2 + 1] = static_cast(unpack_int4_nibble(packed >> 4)); + } } } // namespace @@ -2619,12 +2760,217 @@ void launch_unpack_int4_to_int8_kernel( if (rows == 0 || K_half == 0) return; const int64_t total_packed = rows * K_half; constexpr int block_threads = 256; - const int64_t blocks = (total_packed + block_threads - 1) / block_threads; - unpack_int4_to_int8_kernel<<(blocks), block_threads, 0, stream>>>( + const int64_t blocks = ((total_packed + 7) / 8 + block_threads - 1) / block_threads; + unpack_int4_to_int8_vec8_kernel<<(blocks), block_threads, 0, stream>>>( reinterpret_cast(input), reinterpret_cast(output), - total_packed, - static_cast(K_half)); + total_packed); +} + +void launch_int4_weight_int8_act_gemv_dequant_kernel( + const void* input, + const void* weight, + const void* x_scales, + const void* weight_scales, + const void* bias, + void* output, + int64_t num_rows, + int64_t num_cols, + int64_t K, + int64_t weight_scale_size, + bool has_bias, + int output_dtype_code, + int bias_dtype_code, + cudaStream_t stream) +{ + if (num_cols == 0 || K == 0) return; + if ((K & 3) != 0) { + throw std::runtime_error("int4_weight_int8_act_gemv_dequant requires K divisible by 4"); + } + if (num_cols > static_cast(std::numeric_limits::max()) || + num_rows > static_cast(std::numeric_limits::max()) || + K > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("int4_weight_int8_act_gemv_dequant only supports M,N,K <= INT_MAX"); + } + if (weight_scale_size != 1 && weight_scale_size != num_cols) { + throw std::runtime_error("INT4 packed GEMV weight scale must be scalar or per-output-channel"); + } + + constexpr int kWarpsPerBlock = 8; + const dim3 grid( + static_cast((num_cols + kWarpsPerBlock - 1) / kWarpsPerBlock), + static_cast(num_rows)); + + DISPATCH_FP_DTYPE(output_dtype_code, OutputType, [&] { + if (!has_bias) { + int4_weight_int8_act_gemv_dequant_warp_kernel + <<>>( + static_cast(input), + static_cast(weight), + static_cast(x_scales), + static_cast(weight_scales), + nullptr, + static_cast(output), + static_cast(num_rows), + static_cast(num_cols), + static_cast(K), + static_cast(weight_scale_size), + false); + return; + } + + DISPATCH_FP_DTYPE(bias_dtype_code, BiasType, [&] { + int4_weight_int8_act_gemv_dequant_warp_kernel + <<>>( + static_cast(input), + static_cast(weight), + static_cast(x_scales), + static_cast(weight_scales), + static_cast(bias), + static_cast(output), + static_cast(num_rows), + static_cast(num_cols), + static_cast(K), + static_cast(weight_scale_size), + true); + }); + }); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("CUDA packed INT4 weight GEMV failed: ") + cudaGetErrorString(err)); + } +} + +void launch_int4_weight_int8_act_gemm_dequant_chunked_kernel( + const void* input, + const void* weight, + const void* x_scales, + const void* weight_scales, + const void* bias, + void* output, + void* weight_workspace, + void* acc_workspace, + void* cublas_workspace, + int64_t cublas_workspace_size, + int64_t num_rows, + int64_t num_cols, + int64_t K, + int64_t weight_scale_size, + int64_t chunk_cols, + bool has_bias, + int output_dtype_code, + int bias_dtype_code, + cudaStream_t stream) +{ + if (num_rows == 0 || num_cols == 0 || K == 0) return; + if ((K & 1) != 0) { + throw std::runtime_error("int4_weight_int8_act_gemm_dequant_chunked requires K divisible by 2"); + } + if (chunk_cols <= 0) { + throw std::runtime_error("int4_weight_int8_act_gemm_dequant_chunked requires positive chunk_cols"); + } + if (num_cols > static_cast(std::numeric_limits::max()) || + num_rows > static_cast(std::numeric_limits::max()) || + K > static_cast(std::numeric_limits::max()) || + chunk_cols > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("int4_weight_int8_act_gemm_dequant_chunked only supports M,N,K,chunk <= INT_MAX"); + } + if (weight_scale_size != 1 && weight_scale_size != num_cols) { + throw std::runtime_error("INT4 chunked GEMM weight scale must be scalar or per-output-channel"); + } + + const int64_t K_half = K >> 1; + constexpr int unpack_threads = 256; + constexpr int dequant_threads = 256; + + for (int64_t n0 = 0; n0 < num_cols; n0 += chunk_cols) { + const int64_t cols = std::min(chunk_cols, num_cols - n0); + const int64_t total_packed = cols * K_half; + const int64_t unpack_blocks = ((total_packed + 7) / 8 + unpack_threads - 1) / unpack_threads; + unpack_int4_to_int8_vec8_kernel<<(unpack_blocks), unpack_threads, 0, stream>>>( + static_cast(weight) + n0 * K_half, + static_cast(weight_workspace), + total_packed); + + const float* chunk_weight_scales = static_cast(weight_scales) + (weight_scale_size == 1 ? 0 : n0); + const void* chunk_bias = has_bias ? static_cast(bias) + n0 * fp_dtype_size_bytes(bias_dtype_code) : nullptr; + void* chunk_output = static_cast(output) + n0 * fp_dtype_size_bytes(output_dtype_code); + const bool used_cutlass = ( + num_rows >= 1024 + && (cols >= 4096 || (num_cols == 2560 && cols == 2560)) + && weight_scale_size != 1 + && (!has_bias || bias_dtype_code == 0) + && launch_cutlass_int8_dequant_strided( + input, + weight_workspace, + x_scales, + chunk_weight_scales, + chunk_bias, + chunk_output, + num_rows, + cols, + K, + num_cols, + output_dtype_code, + stream)); + if (used_cutlass) { + continue; + } + + launch_cublas_gemm_int8_kernel( + input, + weight_workspace, + acc_workspace, + num_rows, + cols, + K, + cublas_workspace, + cublas_workspace_size, + stream); + + const int64_t total = num_rows * cols; + const int64_t dequant_blocks = (total + dequant_threads - 1) / dequant_threads; + DISPATCH_FP_DTYPE(output_dtype_code, OutputType, [&] { + if (!has_bias) { + dequantize_int4_weight_int8_act_chunk_kernel + <<(dequant_blocks), dequant_threads, 0, stream>>>( + static_cast(acc_workspace), + static_cast(x_scales), + static_cast(weight_scales), + nullptr, + static_cast(output), + total, + static_cast(num_cols), + static_cast(cols), + static_cast(n0), + static_cast(weight_scale_size), + false); + return; + } + + DISPATCH_FP_DTYPE(bias_dtype_code, BiasType, [&] { + dequantize_int4_weight_int8_act_chunk_kernel + <<(dequant_blocks), dequant_threads, 0, stream>>>( + static_cast(acc_workspace), + static_cast(x_scales), + static_cast(weight_scales), + static_cast(bias), + static_cast(output), + total, + static_cast(num_cols), + static_cast(cols), + static_cast(n0), + static_cast(weight_scale_size), + true); + }); + }); + } + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("CUDA chunked INT4 weight INT8 GEMM failed: ") + cudaGetErrorString(err)); + } } bool launch_cutlass_int4_dequant( diff --git a/comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu b/comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu index b526b96..fc2550d 100644 --- a/comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu +++ b/comfy_kitchen/backends/cuda/ops/cutlass_gemm_int8.cu @@ -73,12 +73,18 @@ struct FusedInt8Gemm { static bool run(const int8_t* A, const int8_t* B, const float* xs, const float* ws, const float* bias, ElementOutput* D, int M, int N, int K, cudaStream_t stream) { + return run_strided(A, B, xs, ws, bias, D, M, N, K, N, stream); + } + + static bool run_strided(const int8_t* A, const int8_t* B, const float* xs, const float* ws, + const float* bias, ElementOutput* D, int M, int N, int K, + int output_stride, cudaStream_t stream) { cutlass::gemm::GemmCoord problem(M, N, K); typename EVTD::Arguments cb{ { { { {}, {const_cast(xs), 0.f, {_1{}, _0{}, M}}, {} }, {const_cast(ws), 0.f, {_0{}, _1{}, N}}, {} }, {const_cast(bias), 0.f, {_0{}, _1{}, N}}, {} }, - {D, {N, _1{}, M * N}} }; + {D, {output_stride, _1{}, M * output_stride}} }; typename Gemm::Arguments args( cutlass::gemm::GemmUniversalMode::kGemm, problem, 1, cb, const_cast(A), const_cast(B), nullptr, nullptr, @@ -131,11 +137,17 @@ struct FusedInt8GemmNoBias { static bool run(const int8_t* A, const int8_t* B, const float* xs, const float* ws, ElementOutput* D, int M, int N, int K, cudaStream_t stream) { + return run_strided(A, B, xs, ws, D, M, N, K, N, stream); + } + + static bool run_strided(const int8_t* A, const int8_t* B, const float* xs, const float* ws, + ElementOutput* D, int M, int N, int K, int output_stride, + cudaStream_t stream) { cutlass::gemm::GemmCoord problem(M, N, K); typename EVTD::Arguments cb{ { { {}, {const_cast(xs), 0.f, {_1{}, _0{}, M}}, {} }, {const_cast(ws), 0.f, {_0{}, _1{}, N}}, {} }, - {D, {N, _1{}, M * N}} }; + {D, {output_stride, _1{}, M * output_stride}} }; typename Gemm::Arguments args( cutlass::gemm::GemmUniversalMode::kGemm, problem, 1, cb, const_cast(A), const_cast(B), nullptr, nullptr, @@ -261,6 +273,92 @@ bool dispatch_fused_no_bias(const int8_t* A, const int8_t* B, const float* xs, c if (best < 0) return false; return runners[best](A, B, xs, ws, D, M, N, K, stream); } + +template +bool dispatch_fused_strided(const int8_t* A, const int8_t* B, const float* xs, const float* ws, + const float* bias, OutT* D, int M, int N, int K, int output_stride, + cudaStream_t stream) { + using Fn = bool (*)(const int8_t*, const int8_t*, const float*, const float*, const float*, OutT*, int, int, int, int, cudaStream_t); + static const Fn runners[] = { + &FusedInt8Gemm::run_strided, + &FusedInt8Gemm::run_strided, + &FusedInt8Gemm::run_strided, + }; + constexpr int NC = sizeof(runners) / sizeof(runners[0]); + + static std::mutex mtx; + static std::map, int> cache; + const std::tuple key{M, N, K}; + + int best; + { + std::lock_guard lk(mtx); + auto it = cache.find(key); + best = (it != cache.end()) ? it->second : -2; + } + if (best == -2) { + best = -1; + float best_ms = 1e30f; + cudaEvent_t s, e; cudaEventCreate(&s); cudaEventCreate(&e); + for (int i = 0; i < NC; ++i) { + if (!runners[i](A, B, xs, ws, bias, D, M, N, K, output_stride, stream)) continue; + cudaStreamSynchronize(stream); + cudaEventRecord(s, stream); + for (int r = 0; r < 3; ++r) runners[i](A, B, xs, ws, bias, D, M, N, K, output_stride, stream); + cudaEventRecord(e, stream); cudaEventSynchronize(e); + float ms = 0.f; cudaEventElapsedTime(&ms, s, e); + if (ms < best_ms) { best_ms = ms; best = i; } + } + cudaEventDestroy(s); cudaEventDestroy(e); + std::lock_guard lk(mtx); + cache[key] = best; + } + if (best < 0) return false; + return runners[best](A, B, xs, ws, bias, D, M, N, K, output_stride, stream); +} + +template +bool dispatch_fused_no_bias_strided(const int8_t* A, const int8_t* B, const float* xs, const float* ws, + OutT* D, int M, int N, int K, int output_stride, + cudaStream_t stream) { + using Fn = bool (*)(const int8_t*, const int8_t*, const float*, const float*, OutT*, int, int, int, int, cudaStream_t); + static const Fn runners[] = { + &FusedInt8GemmNoBias::run_strided, + &FusedInt8GemmNoBias::run_strided, + &FusedInt8GemmNoBias::run_strided, + }; + constexpr int NC = sizeof(runners) / sizeof(runners[0]); + + static std::mutex mtx; + static std::map, int> cache; + const std::tuple key{M, N, K}; + + int best; + { + std::lock_guard lk(mtx); + auto it = cache.find(key); + best = (it != cache.end()) ? it->second : -2; + } + if (best == -2) { + best = -1; + float best_ms = 1e30f; + cudaEvent_t s, e; cudaEventCreate(&s); cudaEventCreate(&e); + for (int i = 0; i < NC; ++i) { + if (!runners[i](A, B, xs, ws, D, M, N, K, output_stride, stream)) continue; + cudaStreamSynchronize(stream); + cudaEventRecord(s, stream); + for (int r = 0; r < 3; ++r) runners[i](A, B, xs, ws, D, M, N, K, output_stride, stream); + cudaEventRecord(e, stream); cudaEventSynchronize(e); + float ms = 0.f; cudaEventElapsedTime(&ms, s, e); + if (ms < best_ms) { best_ms = ms; best = i; } + } + cudaEventDestroy(s); cudaEventDestroy(e); + std::lock_guard lk(mtx); + cache[key] = best; + } + if (best < 0) return false; + return runners[best](A, B, xs, ws, D, M, N, K, output_stride, stream); +} } // namespace extern "C" { @@ -290,6 +388,34 @@ bool launch_cutlass_int8_dequant( default: return false; } } + +bool launch_cutlass_int8_dequant_strided( + const void* A, const void* B, const void* xs, const void* ws, const void* bias, + void* D, int64_t M, int64_t N, int64_t K, int64_t output_stride, int out_dtype_code, + cudaStream_t stream) +{ + if (M == 0 || N == 0 || K == 0) return true; + if (output_stride < N) return false; + const int8_t* a = static_cast(A); + const int8_t* b = static_cast(B); + const float* x = static_cast(xs); + const float* w = static_cast(ws); + const float* bs = static_cast(bias); + if (bs == nullptr) { + switch (out_dtype_code) { + case 0: return dispatch_fused_no_bias_strided(a, b, x, w, static_cast(D), M, N, K, output_stride, stream); + case 1: return dispatch_fused_no_bias_strided(a, b, x, w, static_cast(D), M, N, K, output_stride, stream); + case 2: return dispatch_fused_no_bias_strided(a, b, x, w, static_cast(D), M, N, K, output_stride, stream); + default: return false; + } + } + switch (out_dtype_code) { + case 0: return dispatch_fused_strided(a, b, x, w, bs, static_cast(D), M, N, K, output_stride, stream); + case 1: return dispatch_fused_strided(a, b, x, w, bs, static_cast(D), M, N, K, output_stride, stream); + case 2: return dispatch_fused_strided(a, b, x, w, bs, static_cast(D), M, N, K, output_stride, stream); + default: return false; + } +} } // extern "C" #else // !COMFY_HAVE_CUTLASS -- stub; caller falls back to cuBLAS + separate dequant. @@ -300,4 +426,10 @@ extern "C" bool launch_cutlass_int8_dequant( return false; } +extern "C" bool launch_cutlass_int8_dequant_strided( + const void*, const void*, const void*, const void*, const void*, + void*, int64_t, int64_t, int64_t, int64_t, int, cudaStream_t) { + return false; +} + #endif From ac3d717c884646834b388f73d748387751dac999 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 9 Jul 2026 02:05:56 -0400 Subject: [PATCH 5/7] Support forcing the convrot int4 layers to run the matrix mult in int8. --- comfy_kitchen/backends/cuda/__init__.py | 6 +++++- comfy_kitchen/backends/eager/__init__.py | 1 + comfy_kitchen/backends/eager/convrot_w4a4.py | 3 +++ comfy_kitchen/tensor/convrot_w4a4.py | 17 +++++++++++++++++ tests/test_convrot_w4a4.py | 8 ++++++++ 5 files changed, 34 insertions(+), 1 deletion(-) diff --git a/comfy_kitchen/backends/cuda/__init__.py b/comfy_kitchen/backends/cuda/__init__.py index 441e7ac..a414b6f 100644 --- a/comfy_kitchen/backends/cuda/__init__.py +++ b/comfy_kitchen/backends/cuda/__init__.py @@ -953,8 +953,11 @@ def convrot_w4a4_linear( bias: torch.Tensor | None = None, convrot_groupsize: int = 256, quant_group_size: int = _INT4_GROUP_SIZE, + linear_dtype: str = "int4", ) -> torch.Tensor: """Compute ``x @ W.T + bias`` using ConvRot W4A4 signed INT4 MMA.""" + if linear_dtype not in {"int4", "int8"}: + raise ValueError(f"ConvRot W4A4 linear_dtype must be 'int4' or 'int8', got {linear_dtype!r}") if quant_group_size != _INT4_GROUP_SIZE: raise ValueError(f"int4 MMA kernel requires quant_group_size {_INT4_GROUP_SIZE}") if x.shape[-1] != qweight.shape[-1] * 2: @@ -964,7 +967,7 @@ def convrot_w4a4_linear( orig_shape = x.shape x2d = x.reshape(-1, orig_shape[-1]).contiguous() - if not _cuda_device_supports_native_int4_mma(x2d): + if linear_dtype == "int8" or not _cuda_device_supports_native_int4_mma(x2d): if ( convrot_groupsize == 256 and x2d.shape[-1] % 256 == 0 @@ -2394,6 +2397,7 @@ def _build_constraints() -> dict: "bias": ParamConstraint(dtypes=frozenset({torch.float32, torch.float16, torch.bfloat16})), "convrot_groupsize": ParamConstraint(dtypes=frozenset({int})), "quant_group_size": ParamConstraint(dtypes=frozenset({int})), + "linear_dtype": ParamConstraint(dtypes=frozenset({str})), }, default_devices=cuda_devices, min_compute_capability=(8, 0), diff --git a/comfy_kitchen/backends/eager/__init__.py b/comfy_kitchen/backends/eager/__init__.py index 0850ee0..e689b85 100644 --- a/comfy_kitchen/backends/eager/__init__.py +++ b/comfy_kitchen/backends/eager/__init__.py @@ -248,6 +248,7 @@ def _build_constraints() -> dict: "bias": ParamConstraint(dtypes=standard_floats), "convrot_groupsize": ParamConstraint(dtypes=frozenset({int})), "quant_group_size": ParamConstraint(dtypes=frozenset({int})), + "linear_dtype": ParamConstraint(dtypes=frozenset({str})), }, default_devices=all_devices, ), diff --git a/comfy_kitchen/backends/eager/convrot_w4a4.py b/comfy_kitchen/backends/eager/convrot_w4a4.py index 0e08155..2cf5f30 100644 --- a/comfy_kitchen/backends/eager/convrot_w4a4.py +++ b/comfy_kitchen/backends/eager/convrot_w4a4.py @@ -178,8 +178,11 @@ def convrot_w4a4_linear( bias: torch.Tensor | None = None, convrot_groupsize: int = 256, quant_group_size: int = _INT4_GROUP_SIZE, + linear_dtype: str = "int4", ) -> torch.Tensor: """Compute ``x @ W.T + bias`` using the eager ConvRot W4A4 path.""" + if linear_dtype not in {"int4", "int8"}: + raise ValueError(f"ConvRot W4A4 linear_dtype must be 'int4' or 'int8', got {linear_dtype!r}") if quant_group_size != _INT4_GROUP_SIZE: raise ValueError(f"int4 MMA kernel requires quant_group_size {_INT4_GROUP_SIZE}") if x.shape[-1] != qweight.shape[-1] * 2: diff --git a/comfy_kitchen/tensor/convrot_w4a4.py b/comfy_kitchen/tensor/convrot_w4a4.py index 91a532c..00ac011 100644 --- a/comfy_kitchen/tensor/convrot_w4a4.py +++ b/comfy_kitchen/tensor/convrot_w4a4.py @@ -84,8 +84,11 @@ def convrot_w4a4_linear( bias: torch.Tensor | None = None, convrot_groupsize: int = 256, quant_group_size: int = _INT4_GROUP_SIZE, + linear_dtype: str = "int4", ) -> torch.Tensor: """Compute ``x @ W.T + bias`` using ConvRot W4A4 int4 MMA.""" + if linear_dtype not in {"int4", "int8"}: + raise ValueError(f"ConvRot W4A4 linear_dtype must be 'int4' or 'int8', got {linear_dtype!r}") impl = registry.get_implementation( "convrot_w4a4_linear", kwargs={ @@ -95,6 +98,7 @@ def convrot_w4a4_linear( "bias": bias, "convrot_groupsize": convrot_groupsize, "quant_group_size": quant_group_size, + "linear_dtype": linear_dtype, }, ) return impl( @@ -104,6 +108,7 @@ def convrot_w4a4_linear( bias=bias, convrot_groupsize=convrot_groupsize, quant_group_size=quant_group_size, + linear_dtype=linear_dtype, ) @@ -117,8 +122,14 @@ class TensorCoreConvRotW4A4Layout(QuantizedLayout): class Params(BaseLayoutParams): convrot_groupsize: int = 256 quant_group_size: int = _INT4_GROUP_SIZE + linear_dtype: str = "int4" transposed: bool = False + def __post_init__(self): + super().__post_init__() + if self.linear_dtype not in {"int4", "int8"}: + raise ValueError(f"ConvRot W4A4 linear_dtype must be 'int4' or 'int8', got {self.linear_dtype!r}") + def _tensor_fields(self) -> list[str]: return ["scale"] @@ -132,8 +143,11 @@ def quantize( convrot_groupsize: int = 256, quant_group_size: int = _INT4_GROUP_SIZE, stochastic_rounding: int | None = 0, + linear_dtype: str = "int4", **kwargs, ) -> tuple[torch.Tensor, Params]: + if linear_dtype not in {"int4", "int8"}: + raise ValueError(f"ConvRot W4A4 linear_dtype must be 'int4' or 'int8', got {linear_dtype!r}") qdata, scales = quantize_convrot_w4a4_weight( tensor, convrot_groupsize, @@ -146,6 +160,7 @@ def quantize( orig_shape=tuple(tensor.shape), convrot_groupsize=convrot_groupsize, quant_group_size=quant_group_size, + linear_dtype=linear_dtype, ) return qdata, params @@ -173,6 +188,7 @@ def requantize_kwargs(cls, qtensor: QuantizedTensor) -> dict[str, object]: return { "convrot_groupsize": params.convrot_groupsize, "quant_group_size": params.quant_group_size, + "linear_dtype": params.linear_dtype, } @@ -206,6 +222,7 @@ def _convrot_w4a4_forward(input_tensor: torch.Tensor, weight: QuantizedTensor, b bias=bias, convrot_groupsize=params.convrot_groupsize, quant_group_size=params.quant_group_size, + linear_dtype=params.linear_dtype, ) diff --git a/tests/test_convrot_w4a4.py b/tests/test_convrot_w4a4.py index dd2345e..82a6103 100644 --- a/tests/test_convrot_w4a4.py +++ b/tests/test_convrot_w4a4.py @@ -79,6 +79,14 @@ def test_convrot_w4a4_layout_linear_mm_addmm(seed): assert torch.allclose(out_linear, out_addmm, rtol=1e-5, atol=1e-5) +def test_convrot_w4a4_layout_records_linear_dtype(seed): + w = torch.randn(16, 256, dtype=torch.float32) + qt = QuantizedTensor.from_float(w, "TensorCoreConvRotW4A4Layout", linear_dtype="int8") + + assert qt._params.linear_dtype == "int8" + assert qt._params.convrot_groupsize == 256 + + def test_convrot_w4a4_rejects_bad_groups(seed): w = torch.randn(16, 250, dtype=torch.float32) with pytest.raises(ValueError, match="not divisible by convrot_groupsize"): From 93d19e17155c44b4940b8604b0ba24a9310f10e8 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 9 Jul 2026 15:30:47 -0400 Subject: [PATCH 6/7] Fix issues with large activations. --- comfy_kitchen/backends/cuda/__init__.py | 43 ++++++++++++++++++- .../backends/cuda/ops/convrot_w4a4.cu | 22 ++++++++++ .../backends/cuda/ops/int8_linear.cu | 4 +- 3 files changed, 65 insertions(+), 4 deletions(-) diff --git a/comfy_kitchen/backends/cuda/__init__.py b/comfy_kitchen/backends/cuda/__init__.py index a414b6f..9414e8e 100644 --- a/comfy_kitchen/backends/cuda/__init__.py +++ b/comfy_kitchen/backends/cuda/__init__.py @@ -245,6 +245,37 @@ def _convrot_fused_shared_memory_fits(x: torch.Tensor, k: int, group_size: int) return requested_shared < max_shared +def _convrot_int4_fused_shared_memory_fits(x: torch.Tensor, k: int, group_size: int) -> bool: + if not x.is_cuda or group_size not in (16, 64, 256): + return True + props = torch.cuda.get_device_properties(x.get_device()) + max_shared = getattr(props, "shared_memory_per_block_optin", props.shared_memory_per_block) + if _cuda_device_is_turing(x.get_device()): + max_shared = props.shared_memory_per_block + + if group_size in (16, 64): + block_threads = 512 if k > 4096 else 128 + groups_in_flight = block_threads // (group_size // 4) + requested_shared = (k + groups_in_flight * 2 * group_size) * x.element_size() + return requested_shared < max_shared + + if x.shape[0] != 1 and k <= 4096: + block_threads = 256 + scratch_buffers = 2 + elif x.shape[0] == 1: + block_threads = 512 + scratch_buffers = 2 + elif k == 15360: + block_threads = 640 + scratch_buffers = 1 + else: + block_threads = 1024 + scratch_buffers = 2 + groups_in_flight = block_threads // 64 + requested_shared = (k + groups_in_flight * scratch_buffers * 256) * x.element_size() + return requested_shared < max_shared + + def _should_use_convrot_fused_kernel(x: torch.Tensor, k: int, group_size: int) -> bool: return ( group_size == 256 @@ -899,7 +930,11 @@ def quantize_convrot_w4a4_weight( if weight.shape[-1] % convrot_groupsize != 0: raise ValueError(f"in_features {weight.shape[-1]} not divisible by convrot_groupsize {convrot_groupsize}") weight_2d = weight.contiguous() - if convrot_groupsize in (16, 64, 256) and hasattr(_C, "quantize_int4_rowwise_convrot64"): + if ( + convrot_groupsize in (16, 64, 256) + and hasattr(_C, "quantize_int4_rowwise_convrot64") + and _convrot_int4_fused_shared_memory_fits(weight_2d, weight_2d.shape[-1], convrot_groupsize) + ): qdata, scales = quantize_int4_rowwise_convrot64( weight_2d, convrot_groupsize, @@ -999,7 +1034,11 @@ def convrot_w4a4_linear( x.dtype, ) return out[: x2d.shape[0]].reshape(*orig_shape[:-1], qweight.shape[0]) - if convrot_groupsize in (16, 64, 256) and hasattr(_C, "quantize_int4_rowwise_convrot64"): + if ( + convrot_groupsize in (16, 64, 256) + and hasattr(_C, "quantize_int4_rowwise_convrot64") + and _convrot_int4_fused_shared_memory_fits(x2d, x2d.shape[-1], convrot_groupsize) + ): qact, x_scale = quantize_int4_rowwise_convrot64(x2d, convrot_groupsize) else: h = _build_hadamard(convrot_groupsize, device=x2d.device, dtype=x2d.dtype) diff --git a/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu b/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu index cd24b73..4b68e33 100644 --- a/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu +++ b/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu @@ -1929,6 +1929,11 @@ void launch_quantize_int4_rowwise_kernel( seed); } } + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("CUDA INT4 rowwise quantization failed: ") + cudaGetErrorString(err)); + } } void launch_quantize_int4_rowwise_convrot64_kernel( @@ -1983,6 +1988,13 @@ void launch_quantize_int4_rowwise_convrot64_kernel( static_cast(K), seed); }; + auto check_launch = [] { + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("CUDA INT4 rowwise convrot quantization failed: ") + cudaGetErrorString(err)); + } + }; if (group_size == 16) { const bool use_wide_small_group = K > 4096; if (input_dtype_code == 2) { @@ -2025,6 +2037,7 @@ void launch_quantize_int4_rowwise_convrot64_kernel( launch_small(quantize_int4_rowwise_convrot_small_kernel<16, float, 128, false>, typed_input, 16, 128); } } + check_launch(); return; } if (group_size == 64) { @@ -2069,6 +2082,7 @@ void launch_quantize_int4_rowwise_convrot64_kernel( launch_small(quantize_int4_rowwise_convrot_small_kernel<64, float, 128, false>, typed_input, 64, 128); } } + check_launch(); return; } if (input_dtype_code == 2) { @@ -2153,6 +2167,8 @@ void launch_quantize_int4_rowwise_convrot64_kernel( } } } + + check_launch(); } void launch_quantize_int4_rowwise_convrot64_to_int8_kernel( @@ -2306,6 +2322,12 @@ void launch_quantize_int4_rowwise_convrot64_to_int8_kernel( } } } + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + throw std::runtime_error( + std::string("CUDA INT4 rowwise convrot INT8-output quantization failed: ") + cudaGetErrorString(err)); + } } void launch_dequantize_int4_convrot64_kernel( diff --git a/comfy_kitchen/backends/cuda/ops/int8_linear.cu b/comfy_kitchen/backends/cuda/ops/int8_linear.cu index 20b839b..da66a98 100644 --- a/comfy_kitchen/backends/cuda/ops/int8_linear.cu +++ b/comfy_kitchen/backends/cuda/ops/int8_linear.cu @@ -167,7 +167,7 @@ __global__ void quantize_int8_rowwise_kernel( const int row = static_cast(blockIdx.x); const int tid = threadIdx.x; - const int row_offset = row * K; + const int64_t row_offset = static_cast(row) * K; float abs_max = 0.0f; for (int col = tid; col < K; col += blockDim.x) { @@ -184,7 +184,7 @@ __global__ void quantize_int8_rowwise_kernel( } for (int col = tid; col < K; col += blockDim.x) { - const int idx = row_offset + col; + const int64_t idx = row_offset + col; const float scaled = quant_div_to_float(x[idx], scale); float quantized; if constexpr (STOCHASTIC) { From 314f4bc494554e5e993df1cd361ab7a96d233798 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 9 Jul 2026 17:07:19 -0400 Subject: [PATCH 7/7] Update function used to calculate shared mem usage. --- comfy_kitchen/backends/cuda/__init__.py | 67 ++++++----- tests/test_convrot_w4a4.py | 141 ++++++++++++++++++++++++ 2 files changed, 179 insertions(+), 29 deletions(-) diff --git a/comfy_kitchen/backends/cuda/__init__.py b/comfy_kitchen/backends/cuda/__init__.py index 9414e8e..92a6bbd 100644 --- a/comfy_kitchen/backends/cuda/__init__.py +++ b/comfy_kitchen/backends/cuda/__init__.py @@ -159,7 +159,6 @@ def find_lib_dir(start_dir, lib_pattern): _empty_cuda_tensors: dict[tuple[str, int | None, torch.dtype], torch.Tensor] = {} _turing_device_cache: dict[int, bool] = {} _cutlass_int8_device_cache: dict[int, bool] = {} -_shared_memory_per_block_cache: dict[int, int] = {} _FORCE_INT4_INT8_FALLBACK = os.environ.get("COMFY_KITCHEN_FORCE_INT4_INT8_FALLBACK", "0") == "1" _INT4_PACKED_WEIGHT_SMALL_M_MAX = 8 _INT4_INT8_WEIGHT_CHUNK_N = max(1, int(os.environ.get("COMFY_KITCHEN_INT4_INT8_WEIGHT_CHUNK_N", "4096"))) @@ -227,42 +226,39 @@ def _pad_1d(x: torch.Tensor, padded_size: int) -> torch.Tensor: return torch.cat((x.reshape(-1), padding), dim=0).contiguous() -def _convrot_fused_shared_memory_fits(x: torch.Tensor, k: int, group_size: int) -> bool: - if not x.is_cuda or group_size != 256: - return True +def _max_dynamic_shared_memory_per_block(x: torch.Tensor) -> int: device_index = x.get_device() - max_shared = _shared_memory_per_block_cache.get(device_index) - if max_shared is None: - props = torch.cuda.get_device_properties(device_index) - if _cuda_device_is_turing(device_index): - max_shared = props.shared_memory_per_block - else: - max_shared = getattr(props, "shared_memory_per_block_optin", props.shared_memory_per_block) - _shared_memory_per_block_cache[device_index] = max_shared - # The fused convrot64 kernel stages both fp32 scratch and row data in - # shared memory. Its launch-time request is 8 bytes per column. - requested_shared = k * 8 - return requested_shared < max_shared + props = torch.cuda.get_device_properties(x.get_device()) + if _cuda_device_is_turing(device_index): + return props.shared_memory_per_block + return getattr(props, "shared_memory_per_block_optin", props.shared_memory_per_block) -def _convrot_int4_fused_shared_memory_fits(x: torch.Tensor, k: int, group_size: int) -> bool: - if not x.is_cuda or group_size not in (16, 64, 256): - return True - props = torch.cuda.get_device_properties(x.get_device()) - max_shared = getattr(props, "shared_memory_per_block_optin", props.shared_memory_per_block) - if _cuda_device_is_turing(x.get_device()): - max_shared = props.shared_memory_per_block +def _convrot_int8_fused_shared_memory_bytes(m: int, k: int) -> int: + if m == 1: + block_threads = 512 + elif k == 256: + block_threads = 64 + elif k == 2560: + block_threads = 640 + elif k == 6144: + block_threads = 768 + else: + block_threads = 1024 + groups_in_flight = block_threads // 64 + return (k + groups_in_flight * 2 * 256) * 4 + +def _convrot_int4_fused_shared_memory_bytes(m: int, k: int, group_size: int, dtype_size: int) -> int: if group_size in (16, 64): block_threads = 512 if k > 4096 else 128 groups_in_flight = block_threads // (group_size // 4) - requested_shared = (k + groups_in_flight * 2 * group_size) * x.element_size() - return requested_shared < max_shared + return (k + groups_in_flight * 2 * group_size) * dtype_size - if x.shape[0] != 1 and k <= 4096: + if m != 1 and k <= 4096: block_threads = 256 scratch_buffers = 2 - elif x.shape[0] == 1: + elif m == 1: block_threads = 512 scratch_buffers = 2 elif k == 15360: @@ -272,8 +268,21 @@ def _convrot_int4_fused_shared_memory_fits(x: torch.Tensor, k: int, group_size: block_threads = 1024 scratch_buffers = 2 groups_in_flight = block_threads // 64 - requested_shared = (k + groups_in_flight * scratch_buffers * 256) * x.element_size() - return requested_shared < max_shared + return (k + groups_in_flight * scratch_buffers * 256) * dtype_size + + +def _convrot_fused_shared_memory_fits(x: torch.Tensor, k: int, group_size: int) -> bool: + if not x.is_cuda or group_size != 256: + return True + requested_shared = _convrot_int8_fused_shared_memory_bytes(x.shape[0], k) + return requested_shared < _max_dynamic_shared_memory_per_block(x) + + +def _convrot_int4_fused_shared_memory_fits(x: torch.Tensor, k: int, group_size: int) -> bool: + if not x.is_cuda or group_size not in (16, 64, 256): + return True + requested_shared = _convrot_int4_fused_shared_memory_bytes(x.shape[0], k, group_size, x.element_size()) + return requested_shared < _max_dynamic_shared_memory_per_block(x) def _should_use_convrot_fused_kernel(x: torch.Tensor, k: int, group_size: int) -> bool: diff --git a/tests/test_convrot_w4a4.py b/tests/test_convrot_w4a4.py index 82a6103..b9303b7 100644 --- a/tests/test_convrot_w4a4.py +++ b/tests/test_convrot_w4a4.py @@ -87,6 +87,38 @@ def test_convrot_w4a4_layout_records_linear_dtype(seed): assert qt._params.convrot_groupsize == 256 +@pytest.mark.parametrize( + ("m", "k", "expected"), + [ + (1, 65536, (65536 + 8 * 2 * 256) * 4), + (2, 256, (256 + 1 * 2 * 256) * 4), + (2, 2560, (2560 + 10 * 2 * 256) * 4), + (2, 6144, (6144 + 12 * 2 * 256) * 4), + (2, 15360, (15360 + 16 * 2 * 256) * 4), + ], +) +def test_convrot_int8_fused_shared_memory_bytes(m, k, expected): + assert cuda_backend._convrot_int8_fused_shared_memory_bytes(m, k) == expected + + +@pytest.mark.parametrize( + ("m", "k", "group_size", "dtype_size", "expected"), + [ + (2, 4096, 16, 2, (4096 + 32 * 2 * 16) * 2), + (2, 8192, 16, 2, (8192 + 128 * 2 * 16) * 2), + (2, 4096, 64, 2, (4096 + 8 * 2 * 64) * 2), + (2, 8192, 64, 2, (8192 + 32 * 2 * 64) * 2), + (2, 4096, 256, 2, (4096 + 4 * 2 * 256) * 2), + (1, 65536, 256, 2, (65536 + 8 * 2 * 256) * 2), + (2, 15360, 256, 2, (15360 + 10 * 1 * 256) * 2), + (2, 65536, 256, 2, (65536 + 16 * 2 * 256) * 2), + (2, 65536, 256, 4, (65536 + 16 * 2 * 256) * 4), + ], +) +def test_convrot_int4_fused_shared_memory_bytes(m, k, group_size, dtype_size, expected): + assert cuda_backend._convrot_int4_fused_shared_memory_bytes(m, k, group_size, dtype_size) == expected + + def test_convrot_w4a4_rejects_bad_groups(seed): w = torch.randn(16, 250, dtype=torch.float32) with pytest.raises(ValueError, match="not divisible by convrot_groupsize"): @@ -124,6 +156,115 @@ def test_convrot_w4a4_cuda_no_bias_large_m(seed): assert out.dtype == torch.bfloat16 +def test_convrot_cuda_shared_memory_fit_matches_device_limit(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + x_int8 = torch.empty((2, 65536), device="cuda", dtype=torch.float16) + max_shared = cuda_backend._max_dynamic_shared_memory_per_block(x_int8) + int8_requested = cuda_backend._convrot_int8_fused_shared_memory_bytes(x_int8.shape[0], x_int8.shape[1]) + + assert cuda_backend._convrot_fused_shared_memory_fits(x_int8, x_int8.shape[1], 256) == ( + int8_requested < max_shared + ) + + x_int4 = torch.empty((2, 65536), device="cuda", dtype=torch.float16) + int4_requested = cuda_backend._convrot_int4_fused_shared_memory_bytes( + x_int4.shape[0], + x_int4.shape[1], + 256, + x_int4.element_size(), + ) + + assert cuda_backend._convrot_int4_fused_shared_memory_fits(x_int4, x_int4.shape[1], 256) == ( + int4_requested < max_shared + ) + + +@pytest.mark.parametrize("linear_dtype", ["int4", "int8"]) +@pytest.mark.parametrize("convrot_groupsize", [16, 64, 256]) +def test_convrot_w4a4_cuda_linear_handles_large_fp16_activations(seed, linear_dtype, convrot_groupsize): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + torch.manual_seed(789) + x = torch.empty(8, 256, device="cuda", dtype=torch.float16) + pattern = torch.tensor([65504.0, 65504.0, 65504.0, -65504.0], device="cuda", dtype=torch.float16) + x.copy_(pattern.repeat(x.numel() // pattern.numel()).reshape_as(x)) + w = (torch.randn(32, 256, device="cuda", dtype=torch.float16) * 1.0e-4).contiguous() + bias = torch.zeros(32, device="cuda", dtype=torch.float16) + qt = QuantizedTensor.from_float( + w, + "TensorCoreConvRotW4A4Layout", + convrot_groupsize=convrot_groupsize, + linear_dtype=linear_dtype, + ) + + out = functional.linear(x, qt, bias) + qact, scale = cuda_backend.quantize_int4_rowwise_convrot64(x, convrot_groupsize) + q_values = _unpack_int4_row_major(qact) + + assert out.shape == (8, 32) + assert out.dtype == torch.float16 + assert torch.isfinite(out).all() + assert torch.isfinite(scale).all() + assert q_values.min().item() >= -7 + assert q_values.max().item() <= 7 + + +@pytest.mark.parametrize("linear_dtype", ["int4", "int8"]) +@pytest.mark.parametrize(("m", "k"), [(128, 15360), (1152, 6912)]) +def test_convrot_w4a4_cuda_linear_handles_large_fp16_activation_shapes(seed, linear_dtype, m, k): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + torch.manual_seed(790) + x = torch.empty(m, k, device="cuda", dtype=torch.float16) + pattern = torch.tensor([65504.0, 65504.0, 65504.0, -65504.0], device="cuda", dtype=torch.float16) + x.copy_(pattern.repeat(x.numel() // pattern.numel()).reshape_as(x)) + w = (torch.randn(64, k, device="cuda", dtype=torch.float16) * 1.0e-5).contiguous() + qt = QuantizedTensor.from_float( + w, + "TensorCoreConvRotW4A4Layout", + convrot_groupsize=256, + linear_dtype=linear_dtype, + ) + + out = functional.linear(x, qt) + qact, scale = cuda_backend.quantize_int4_rowwise_convrot64(x, 256) + q_values = _unpack_int4_row_major(qact) + + assert out.shape == (m, 64) + assert out.dtype == torch.float16 + assert torch.isfinite(out).all() + assert torch.isfinite(scale).all() + assert q_values.min().item() >= -7 + assert q_values.max().item() <= 7 + + +@pytest.mark.parametrize("linear_dtype", ["int4", "int8"]) +@pytest.mark.parametrize(("m", "k", "n"), [(4192, 6144, 1536), (4192, 16384, 512)]) +def test_convrot_w4a4_cuda_linear_handles_big_activation_tensors(seed, linear_dtype, m, k, n): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + torch.manual_seed(791) + x = torch.randn(m, k, device="cuda", dtype=torch.float16) + w = torch.randn(n, k, device="cuda", dtype=torch.float16) + qt = QuantizedTensor.from_float( + w, + "TensorCoreConvRotW4A4Layout", + convrot_groupsize=256, + linear_dtype=linear_dtype, + ) + + out = functional.linear(x, qt) + + assert out.shape == (m, n) + assert out.dtype == torch.float16 + assert torch.isfinite(out).all() + + def test_convrot_w4a4_cuda_large_k_quantize_matches_reference(seed): if not torch.cuda.is_available(): pytest.skip("CUDA required")