diff --git a/comfy_kitchen/backends/hip/__init__.py b/comfy_kitchen/backends/hip/__init__.py index 402b3e99..35324057 100644 --- a/comfy_kitchen/backends/hip/__init__.py +++ b/comfy_kitchen/backends/hip/__init__.py @@ -1246,8 +1246,8 @@ def quantize_svdquant_w4a4( r = lora_down.shape[1] m_pad = -(-m // pad_size) * pad_size - # The kernels take M, K and R with no bounds of their own, so every operand - # has to match those extents and sit on x's device. + # The quantizer takes raw pointers, so its operands must be contiguous and + # live on x's device. xc = x.contiguous() # The kernel decodes smooth with the same dtype code it uses for ascales, which # is allocated from x.dtype, so a smooth of any other dtype would be read as @@ -1266,9 +1266,12 @@ def quantize_svdquant_w4a4( _dl(xc), _dl(smooth), _dl(q), _dl(ascales), m, m_pad, k, act_unsigned, _stream(x), ) - _C.svdquant_lora_down( - _dl(lora_src), _dl(lora_down), _dl(lora_act[:m]), m, k, r, _stream(x) - ) + if m > 0: + lora_act_rows = lora_act[:m] + if lora_act_rows.dtype == lora_src.dtype and lora_act_rows.is_contiguous(): + torch.mm(lora_src, lora_down, out=lora_act_rows) + else: + lora_act_rows.copy_(lora_src @ lora_down, non_blocking=True) return q, ascales, lora_act diff --git a/comfy_kitchen/backends/hip/dlpack_bindings.cpp b/comfy_kitchen/backends/hip/dlpack_bindings.cpp index b7e83084..efe272e4 100644 --- a/comfy_kitchen/backends/hip/dlpack_bindings.cpp +++ b/comfy_kitchen/backends/hip/dlpack_bindings.cpp @@ -83,8 +83,6 @@ void launch_adaln_kernel(const void*, const void*, const void*, void*, int, int, int, int, int, bool, hipStream_t); void launch_gemv_awq_kernel(const void*, const void*, const void*, const void*, const void*, void*, int, int, int, int, int, int, int, int, hipStream_t); -void launch_svdquant_lora_down_kernel(const void*, const void*, void*, int, int, int, int, int, - hipStream_t); void launch_svdquant_quant_kernel(const void*, const void*, void*, void*, int, int, int, int, int, bool, hipStream_t); void launch_svdquant_gemm_kernel(const void*, const void*, void*, const void*, const void*, @@ -947,27 +945,6 @@ void gemv_awq_w4a16(nb::ndarray<> x, nb::ndarray<> qweight, nb::ndarray<> wscale check_hip_launch(); } -void svdquant_lora_down(nb::ndarray<> x, nb::ndarray<> lora_down, nb::ndarray<> lora_act, int M, - int K, int R, uintptr_t stream_ptr) { - constexpr const char* kFn = "svdquant_lora_down"; - require_nonneg(M, kFn, "M"); - require_nonneg(K, kFn, "K"); - require_nonneg(R, kFn, "R"); - require_dtype(x, 0, 2, kFn, "x"); - require_dtype(lora_down, 0, 2, kFn, "lora_down"); - // The launcher writes lora_act through a float*, so it must be float32 storage. - require_dtype(lora_act, 0, 0, kFn, "lora_act"); - require_len(x, static_cast(M) * K, kFn, "x"); - require_len(lora_down, static_cast(K) * R, kFn, "lora_down"); - require_len(lora_act, static_cast(M) * R, kFn, "lora_act"); - - launch_svdquant_lora_down_kernel(x.data(), lora_down.data(), lora_act.data(), M, K, R, - map_dtype_to_code(x.dtype()), - map_dtype_to_code(lora_down.dtype()), - reinterpret_cast(stream_ptr)); - check_hip_launch(); -} - void svdquant_quantize(nb::ndarray<> x, nb::ndarray<> smooth, nb::ndarray<> q, nb::ndarray<> ascales, int M, int M_pad, int K, bool act_unsigned, uintptr_t stream_ptr) { @@ -1539,7 +1516,6 @@ NB_MODULE(_C, m) { m.def("apply_rope", &apply_rope); m.def("rms_rope", &rms_rope); m.def("gemv_awq_w4a16", &gemv_awq_w4a16); - m.def("svdquant_lora_down", &svdquant_lora_down); m.def("svdquant_quantize", &svdquant_quantize); m.def("svdquant_gemm", &svdquant_gemm); } diff --git a/comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip b/comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip index 107e46c0..7c27b76f 100644 --- a/comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip +++ b/comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip @@ -26,28 +26,9 @@ namespace comfy::hip_backend { constexpr int kSvdGroup = 64; // quantization group, in elements constexpr int kSvdGroupBytes = 32; // ... and in packed int4 bytes -// x @ proj_down -> (M, R), fp32. One thread per output column: proj_down is -// (K, R), so consecutive threads read consecutive addresses. -__global__ __launch_bounds__(256) void lora_down_kernel( - const void* __restrict__ x, const void* __restrict__ lora_down, - float* __restrict__ lora_act, int M, int K, int R, int x_code, int d_code) { - - const int m = blockIdx.x; - const int r = threadIdx.x; - if (r >= R) return; - - float sum = 0.0f; - for (int k = 0; k < K; ++k) { - sum += load_in(x, static_cast(m) * K + k, x_code) * - load_in(lora_down, static_cast(k) * R + r, d_code); - } - lora_act[static_cast(m) * R + r] = sum; -} - -// Smooth, then per-row per-group int4 quantize + pack. ascales is stored -// transposed as (K/64, M_pad). +// Scalar fallback for operands with different dtypes. template -__global__ __launch_bounds__(256) void svdquant_quant_kernel( +__global__ __launch_bounds__(256) void svdquant_quant_scalar_kernel( const void* __restrict__ x, const void* __restrict__ smooth, int8_t* __restrict__ q, void* __restrict__ ascales, int M, int M_pad, int K, int x_code, int s_code) { @@ -91,6 +72,100 @@ __global__ __launch_bounds__(256) void svdquant_quant_kernel( } } +// Smooth, then per-row per-group int4 quantize + pack. Sixteen threads process +// one 64-element group, keeping four smoothed values per thread so x and smooth +// are read once. ascales is stored transposed as (K/64, M_pad). +template +__global__ __launch_bounds__(256) void svdquant_quant_kernel( + const Elem* __restrict__ x, const Elem* __restrict__ smooth, + int8_t* __restrict__ q, Elem* __restrict__ ascales, + int M, int M_pad, int K) { + + constexpr int kThreadsPerGroup = 16; + constexpr int kElemsPerThread = kSvdGroup / kThreadsPerGroup; + constexpr int kGroupsPerBlock = 256 / kThreadsPerGroup; + constexpr float kQMax = UNSIGNED ? 15.0f : 7.0f; + constexpr int kQMin = UNSIGNED ? 0 : -7; + constexpr int kQMaxInt = UNSIGNED ? 15 : 7; + + const int m = blockIdx.x; + const int lane = threadIdx.x % kThreadsPerGroup; + const int group_slot = threadIdx.x / kThreadsPerGroup; + const int ngroups = K / kSvdGroup; + + for (int g = group_slot; g < ngroups; g += kGroupsPerBlock) { + const int elem = lane * kElemsPerThread; + const int kbase = g * kSvdGroup + elem; + const int64_t xbase = static_cast(m) * K + kbase; + + float values[kElemsPerThread]; + float absmax = 0.0f; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) { + values[i] = static_cast(x[xbase + i]) / + static_cast(smooth[kbase + i]); + absmax = fmaxf(absmax, fabsf(values[i])); + } + +#pragma unroll + for (int offset = kThreadsPerGroup / 2; offset > 0; offset >>= 1) { + absmax = fmaxf(absmax, __shfl_xor(absmax, offset, kThreadsPerGroup)); + } + + absmax = fmaxf(absmax, 1e-10f); + const float scale = absmax / kQMax; + const float inv = kQMax / absmax; + if (lane == 0) { + ascales[static_cast(g) * M_pad + m] = static_cast(scale); + } + + int quantized[kElemsPerThread]; +#pragma unroll + for (int i = 0; i < kElemsPerThread; ++i) { + int v = static_cast(rintf(values[i] * inv)); + quantized[i] = v < kQMin ? kQMin : (v > kQMaxInt ? kQMaxInt : v); + } + + const uint16_t packed = + static_cast((quantized[0] & 0xF) | + ((quantized[1] & 0xF) << 4) | + ((quantized[2] & 0xF) << 8) | + ((quantized[3] & 0xF) << 12)); + int8_t* qrow = q + static_cast(m) * (K / 2) + g * kSvdGroupBytes; + reinterpret_cast(qrow)[lane] = packed; + } +} + +template +void launch_svdquant_quant( + const void* x, const void* smooth, int8_t* q, void* ascales, + int M, int M_pad, int K, int x_code, int s_code, hipStream_t stream) { + + if (x_code == s_code) { + if (x_code == 0) { + svdquant_quant_kernel<<>>( + static_cast(x), static_cast(smooth), q, + static_cast(ascales), M, M_pad, K); + return; + } + if (x_code == 1) { + svdquant_quant_kernel<__half, UNSIGNED><<>>( + static_cast(x), static_cast(smooth), q, + static_cast<__half*>(ascales), M, M_pad, K); + return; + } + if (x_code == 2) { + svdquant_quant_kernel<__bf16, UNSIGNED><<>>( + static_cast(x), static_cast(smooth), q, + static_cast<__bf16*>(ascales), M, M_pad, K); + return; + } + } + + svdquant_quant_scalar_kernel<<>>( + x, smooth, q, ascales, M, M_pad, K, x_code, s_code); +} + // Group-scaled int4 GEMM with the LoRA-up correction folded into the writeback. template __global__ __launch_bounds__(256) void svdquant_gemm_kernel( @@ -236,19 +311,6 @@ __global__ __launch_bounds__(256) void svdquant_gemm_kernel( } // namespace comfy::hip_backend -extern "C" void launch_svdquant_lora_down_kernel( - const void* x, const void* lora_down, void* lora_act, int M, int K, int R, int x_code, - int d_code, hipStream_t stream) { - - using namespace comfy::hip_backend; - if (R > 256) throw std::runtime_error("svdquant: LoRA rank above 256 is not supported"); - if (M == 0) { - return; // a zero-block launch is hipErrorInvalidConfiguration - } - lora_down_kernel<<>>(x, lora_down, static_cast(lora_act), M, K, R, - x_code, d_code); -} - extern "C" void launch_svdquant_quant_kernel( const void* x, const void* smooth, void* q, void* ascales, int M, int M_pad, int K, int x_code, int s_code, bool act_unsigned, hipStream_t stream) { @@ -265,11 +327,13 @@ extern "C" void launch_svdquant_quant_kernel( " must be a multiple of 64"); } if (act_unsigned) { - svdquant_quant_kernel<<>>( - x, smooth, static_cast(q), ascales, M, M_pad, K, x_code, s_code); + launch_svdquant_quant( + x, smooth, static_cast(q), ascales, + M, M_pad, K, x_code, s_code, stream); } else { - svdquant_quant_kernel<<>>( - x, smooth, static_cast(q), ascales, M, M_pad, K, x_code, s_code); + launch_svdquant_quant( + x, smooth, static_cast(q), ascales, + M, M_pad, K, x_code, s_code, stream); } } diff --git a/tests/test_hip_wmma.py b/tests/test_hip_wmma.py index 02790ab5..d3fd6089 100644 --- a/tests/test_hip_wmma.py +++ b/tests/test_hip_wmma.py @@ -1975,6 +1975,51 @@ def test_svdquant_validates_its_operands(hip): assert torch.isfinite(out).all() +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_svdquant_lora_down_matches_torch_mm(hip, dtype): + """LoRA-down uses the backend matmul and preserves padded zero rows.""" + torch.manual_seed(0) + m, k, rank = 19, 128, 17 + x = torch.randn(m, k, device=DEV, dtype=dtype) + smooth = torch.ones(k, device=DEV, dtype=dtype) + lora_down = torch.randn(k, rank, device=DEV, dtype=dtype) + + _, _, lora_act = hip.quantize_svdquant_w4a4( + x, smooth, lora_down, pad_size=16 + ) + ref = torch.mm(x, lora_down).float() + + torch.testing.assert_close(lora_act[:m], ref) + assert torch.count_nonzero(lora_act[m:]) == 0 + + +@pytest.mark.parametrize("k", [128, 1088]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("act_unsigned", [False, True]) +def test_svdquant_activation_quantizer_group_layout(hip, dtype, act_unsigned, k): + """Each 16-thread group preserves nibble order and transposed scales.""" + from comfy_kitchen.backends.eager.svdquant import _pack_int4_row_major + + m = 3 + values = torch.arange(k, device=DEV).repeat(m, 1) + if act_unsigned: + values = values.remainder(16) + else: + values = values.remainder(15) - 7 + x = values.to(dtype) + smooth = torch.ones(k, device=DEV, dtype=dtype) + lora_down = torch.zeros(k, 1, device=DEV, dtype=dtype) + + q, ascales, _ = hip.quantize_svdquant_w4a4( + x, smooth, lora_down, pad_size=16, act_unsigned=act_unsigned + ) + + assert torch.equal(q[:m], _pack_int4_row_major(values)) + assert torch.count_nonzero(q[m:]) == 0 + assert torch.equal(ascales[:, :m], torch.ones_like(ascales[:, :m])) + assert torch.count_nonzero(ascales[:, m:]) == 0 + + def test_stochastic_rounding_rejects_non_contiguous_rng(hip): """The kernel writes the result into rng, which a copy would silently discard.""" x = torch.randn(8, 16, device=DEV, dtype=torch.bfloat16)