diff --git a/comfy_kitchen/__init__.py b/comfy_kitchen/__init__.py index a8172ef8..96514d95 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 9f2b1db8..8beae7a4 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 228ec9c8..92a6bbd5 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, @@ -147,7 +159,9 @@ 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"))) def _cuda_device_is_turing(device_index: int) -> bool: @@ -172,6 +186,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 @@ -202,22 +226,63 @@ def _pad_1d(x: torch.Tensor, padded_size: int) -> torch.Tensor: return torch.cat((x.reshape(-1), padding), dim=0).contiguous() +def _max_dynamic_shared_memory_per_block(x: torch.Tensor) -> int: + device_index = x.get_device() + 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_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) + return (k + groups_in_flight * 2 * group_size) * dtype_size + + if m != 1 and k <= 4096: + block_threads = 256 + scratch_buffers = 2 + elif m == 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 + 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 - 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 + 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: @@ -487,6 +552,518 @@ 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_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, + 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() + + 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 prefer_cublas_fallback + and not _DISABLE_CUTLASS_INT8 + and _cuda_device_supports_cutlass_int8_dequant(x_int8) + ): + 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) + _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 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: + 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) + + +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") + 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, + 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, + 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: + 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 linear_dtype == "int8" or not _cuda_device_supports_native_int4_mma(x2d): + 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) + 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, + qweight, + 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") + 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) + 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 +2398,67 @@ 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})), + "linear_dtype": ParamConstraint(dtypes=frozenset({str})), + }, + 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 b788ce94..ce65aa95 100644 --- a/comfy_kitchen/backends/cuda/dlpack_bindings.cpp +++ b/comfy_kitchen/backends/cuda/dlpack_bindings.cpp @@ -750,6 +750,111 @@ 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); + + 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, @@ -763,6 +868,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 +1056,377 @@ 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); +} + +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( @@ -960,6 +1449,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 +1964,86 @@ 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("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"), @@ -1460,6 +2055,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 00000000..4b68e332 --- /dev/null +++ b/comfy_kitchen/backends/cuda/ops/convrot_w4a4.cu @@ -0,0 +1,3038 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// 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 +#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 + +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; +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)); +} + +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()); +} + +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; +} + +#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( + finite_absmax_for_int4_scale(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( + finite_absmax_for_int4_scale(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( + finite_absmax_for_int4_scale(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( + finite_absmax_for_int4_scale(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); + } + } + } +} + +__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) +{ + 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; + } + + 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 + +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); + } + } + + 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( + 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); + }; + 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) { + 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); + } + } + check_launch(); + 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); + } + } + check_launch(); + 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); + } + } + } + + check_launch(); +} + +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); + } + } + } + + 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( + 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 + 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); +} + +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( + 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 5bb69562..fc2550d4 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" { @@ -275,7 +373,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); @@ -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 diff --git a/comfy_kitchen/backends/cuda/ops/int8_linear.cu b/comfy_kitchen/backends/cuda/ops/int8_linear.cu index 20b839b6..da66a98e 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) { diff --git a/comfy_kitchen/backends/eager/__init__.py b/comfy_kitchen/backends/eager/__init__.py index 02002962..e689b85d 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,43 @@ 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})), + "linear_dtype": ParamConstraint(dtypes=frozenset({str})), + }, + 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 00000000..2cf5f307 --- /dev/null +++ b/comfy_kitchen/backends/eager/convrot_w4a4.py @@ -0,0 +1,199 @@ +# 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, + 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: + 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 e1a9bf56..aa0efd8c 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 00000000..00ac0112 --- /dev/null +++ b/comfy_kitchen/tensor/convrot_w4a4.py @@ -0,0 +1,261 @@ +# 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, + 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={ + "x": x, + "qweight": qweight, + "wscales": wscales, + "bias": bias, + "convrot_groupsize": convrot_groupsize, + "quant_group_size": quant_group_size, + "linear_dtype": linear_dtype, + }, + ) + return impl( + x, + qweight, + wscales, + bias=bias, + convrot_groupsize=convrot_groupsize, + quant_group_size=quant_group_size, + linear_dtype=linear_dtype, + ) + + +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 + 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"] + + 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, + 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, + 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, + linear_dtype=linear_dtype, + ) + 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, + "linear_dtype": params.linear_dtype, + } + + +@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, + linear_dtype=params.linear_dtype, + ) + + +@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) diff --git a/tests/test_convrot_w4a4.py b/tests/test_convrot_w4a4.py new file mode 100644 index 00000000..b9303b79 --- /dev/null +++ b/tests/test_convrot_w4a4.py @@ -0,0 +1,320 @@ +"""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_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 + + +@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"): + 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_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") + + 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)