From 2bf7b4b3952b1ef0e1eeee0f927e6c580c496b4c Mon Sep 17 00:00:00 2001 From: tangzzycc <3081129260@qq.com> Date: Sat, 22 Aug 2026 13:41:01 +0800 Subject: [PATCH 1/3] Optimize HIP SVDQuant INT4 quantization and WMMA LoRA-down Use 16 threads per activation quantization group to reuse loaded values, and add FP16/BF16 WMMA LoRA-down while preserving scalar fallbacks. --- comfy_kitchen/backends/hip/__init__.py | 6 +- .../backends/hip/dlpack_bindings.cpp | 32 ++- .../backends/hip/ops/svdquant_w4a4.hip | 230 +++++++++++++++++- tests/test_hip_wmma.py | 66 +++++ 4 files changed, 313 insertions(+), 21 deletions(-) diff --git a/comfy_kitchen/backends/hip/__init__.py b/comfy_kitchen/backends/hip/__init__.py index 402b3e99..f3ada51e 100644 --- a/comfy_kitchen/backends/hip/__init__.py +++ b/comfy_kitchen/backends/hip/__init__.py @@ -1266,8 +1266,10 @@ 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) + lora_down_op = _C.svdquant_lora_down_wmma if has_wmma() else _C.svdquant_lora_down + lora_down_op( + _dl(lora_src), _dl(lora_down), _dl(lora_act[:m]), + m, k, r, _stream(x), ) 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..de8c539f 100644 --- a/comfy_kitchen/backends/hip/dlpack_bindings.cpp +++ b/comfy_kitchen/backends/hip/dlpack_bindings.cpp @@ -85,6 +85,8 @@ void launch_gemv_awq_kernel(const void*, const void*, const void*, const 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_lora_down_wmma_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,8 +949,8 @@ 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) { +void svdquant_lora_down_impl(nb::ndarray<> x, nb::ndarray<> lora_down, nb::ndarray<> lora_act, + int M, int K, int R, uintptr_t stream_ptr, bool use_wmma) { constexpr const char* kFn = "svdquant_lora_down"; require_nonneg(M, kFn, "M"); require_nonneg(K, kFn, "K"); @@ -961,13 +963,30 @@ void svdquant_lora_down(nb::ndarray<> x, nb::ndarray<> lora_down, nb::ndarray<> 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)); + const int x_code = map_dtype_to_code(x.dtype()); + const int d_code = map_dtype_to_code(lora_down.dtype()); + if (use_wmma) { + launch_svdquant_lora_down_wmma_kernel( + x.data(), lora_down.data(), lora_act.data(), M, K, R, x_code, d_code, + reinterpret_cast(stream_ptr)); + } else { + launch_svdquant_lora_down_kernel( + x.data(), lora_down.data(), lora_act.data(), M, K, R, x_code, d_code, + reinterpret_cast(stream_ptr)); + } 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) { + svdquant_lora_down_impl(x, lora_down, lora_act, M, K, R, stream_ptr, false); +} + +void svdquant_lora_down_wmma(nb::ndarray<> x, nb::ndarray<> lora_down, nb::ndarray<> lora_act, + int M, int K, int R, uintptr_t stream_ptr) { + svdquant_lora_down_impl(x, lora_down, lora_act, M, K, R, stream_ptr, true); +} + 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) { @@ -1540,6 +1559,7 @@ NB_MODULE(_C, m) { 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_lora_down_wmma", &svdquant_lora_down_wmma); 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..851ade36 100644 --- a/comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip +++ b/comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip @@ -26,9 +26,8 @@ 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( +// Scalar fallback for FP32 and mixed input dtypes. +__global__ __launch_bounds__(256) void lora_down_scalar_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) { @@ -44,10 +43,98 @@ __global__ __launch_bounds__(256) void lora_down_kernel( 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). +// x[M, K] @ lora_down[K, R] -> lora_act[M, R], with fp32 accumulation. +// +// One wave computes a 16-row by (TN * 16)-column output tile. The A fragment is +// contiguous in x. lora_down is row-major KxR rather than the transposed RxK +// layout consumed by the WMMA instruction, so each lane gathers one output +// column across the K step. For a fixed K element the 16 fragment rows are +// adjacent in memory, keeping those gathers coalesced without an LDS transpose. +template +__global__ __launch_bounds__(kWave) void lora_down_wmma_kernel( + const typename Mma::Elem* __restrict__ x, + const typename Mma::Elem* __restrict__ lora_down, + float* __restrict__ lora_act, int M, int K, int R) { + + const int lane = threadIdx.x; + const int frag_r = frag_row(lane); + const int m0 = blockIdx.y * 16; + const int r0 = blockIdx.x * (TN * 16); + + typename Mma::Acc acc[TN]; +#pragma unroll + for (int n = 0; n < TN; ++n) { + acc[n] = Mma::zero(); + } + + for (int k0 = 0; k0 < K; k0 += 16) { + typename Mma::Frag a{}; + const int row = m0 + frag_r; + if (row < M) { + a = load_frag_16bit( + x + static_cast(row) * K + k0, lane); + } + +#pragma unroll + for (int n = 0; n < TN; ++n) { + typename Mma::Frag b{}; + const int col = r0 + n * 16 + frag_r; + if (col < R) { + const int kbase = k0 + Mma::frag_base(lane); +#pragma unroll + for (int i = 0; i < Mma::kFragElems; ++i) { + b[i] = lora_down[static_cast(kbase + i) * R + col]; + } + } + acc[n] = Mma::mma(a, b, acc[n]); + } + } + + const int col_lane = acc_col(lane); +#pragma unroll + for (int e = 0; e < 8; ++e) { + const int row = m0 + acc_row(lane, e); + if (row >= M) continue; +#pragma unroll + for (int n = 0; n < TN; ++n) { + const int col = r0 + n * 16 + col_lane; + if (col < R) { + lora_act[static_cast(row) * R + col] = Mma::get(acc[n], e); + } + } + } +} + +template +void launch_lora_down_wmma( + const void* x, const void* lora_down, float* lora_act, + int M, int K, int R, hipStream_t stream) { + + using Elem = typename Mma::Elem; + if (R <= 16) { + constexpr int TN = 1; + const dim3 grid((R + TN * 16 - 1) / (TN * 16), (M + 15) / 16); + lora_down_wmma_kernel<<>>( + static_cast(x), static_cast(lora_down), + lora_act, M, K, R); + } else if (R <= 32) { + constexpr int TN = 2; + const dim3 grid((R + TN * 16 - 1) / (TN * 16), (M + 15) / 16); + lora_down_wmma_kernel<<>>( + static_cast(x), static_cast(lora_down), + lora_act, M, K, R); + } else { + constexpr int TN = 4; + const dim3 grid((R + TN * 16 - 1) / (TN * 16), (M + 15) / 16); + lora_down_wmma_kernel<<>>( + static_cast(x), static_cast(lora_down), + lora_act, M, K, R); + } +} + +// 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 +178,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( @@ -242,11 +423,32 @@ extern "C" void launch_svdquant_lora_down_kernel( using namespace comfy::hip_backend; if (R > 256) throw std::runtime_error("svdquant: LoRA rank above 256 is not supported"); - if (M == 0) { + if (M == 0 || R == 0) { + return; // a zero-block launch is hipErrorInvalidConfiguration + } + lora_down_scalar_kernel<<>>( + x, lora_down, static_cast(lora_act), M, K, R, x_code, d_code); +} + +extern "C" void launch_svdquant_lora_down_wmma_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 || R == 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); + if ((K & 15) == 0 && x_code == 1 && d_code == 1) { + launch_lora_down_wmma( + x, lora_down, static_cast(lora_act), M, K, R, stream); + } else if ((K & 15) == 0 && x_code == 2 && d_code == 2) { + launch_lora_down_wmma( + x, lora_down, static_cast(lora_act), M, K, R, stream); + } else { + launch_svdquant_lora_down_kernel( + x, lora_down, lora_act, M, K, R, x_code, d_code, stream); + } } extern "C" void launch_svdquant_quant_kernel( @@ -265,11 +467,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..df1a9d32 100644 --- a/tests/test_hip_wmma.py +++ b/tests/test_hip_wmma.py @@ -1975,6 +1975,72 @@ def test_svdquant_validates_its_operands(hip): assert torch.isfinite(out).all() +@needs_wmma +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("rank", [1, 7, 16, 17, 32, 33, 64, 65, 96, 128, 192, 256]) +def test_svdquant_lora_down_matches_fp32_reference(hip, dtype, rank): + """The WMMA LoRA-down path handles partial M/R tiles and accumulates in fp32.""" + torch.manual_seed(rank) + m, k = 19, 128 + 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 = x.float() @ lora_down.float() + + torch.testing.assert_close( + lora_act[:m], ref, rtol=2e-3, atol=2e-3 + ) + assert torch.count_nonzero(lora_act[m:]) == 0 + + +def test_svdquant_lora_down_scalar_fallback(hip, monkeypatch): + """Devices without matrix cores retain the scalar LoRA-down path.""" + torch.manual_seed(0) + m, k, rank = 5, 64, 7 + x = torch.randn(m, k, device=DEV, dtype=torch.float16) + smooth = torch.ones(k, device=DEV, dtype=torch.float16) + lora_down = torch.randn(k, rank, device=DEV, dtype=torch.float16) + monkeypatch.setattr(hip, "has_wmma", lambda: False) + + _, _, lora_act = hip.quantize_svdquant_w4a4( + x, smooth, lora_down, pad_size=16 + ) + ref = x.float() @ lora_down.float() + + torch.testing.assert_close(lora_act[:m], ref, rtol=2e-3, atol=2e-3) + assert torch.count_nonzero(lora_act[m:]) == 0 + + +@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): + """Each 16-thread group preserves nibble order and transposed scales.""" + from comfy_kitchen.backends.eager.svdquant import _pack_int4_row_major + + m, k = 3, 128 + 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) From 4ff75eb82e52eeb1e3f6055b3c94c60282080e8e Mon Sep 17 00:00:00 2001 From: tangzzycc <3081129260@qq.com> Date: Sat, 22 Aug 2026 14:05:50 +0800 Subject: [PATCH 2/3] Address SVDQuant review feedback --- .../backends/hip/dlpack_bindings.cpp | 12 +++--- tests/test_hip_wmma.py | 37 ++++++++++++++++++- 2 files changed, 42 insertions(+), 7 deletions(-) diff --git a/comfy_kitchen/backends/hip/dlpack_bindings.cpp b/comfy_kitchen/backends/hip/dlpack_bindings.cpp index de8c539f..ff90c281 100644 --- a/comfy_kitchen/backends/hip/dlpack_bindings.cpp +++ b/comfy_kitchen/backends/hip/dlpack_bindings.cpp @@ -949,9 +949,9 @@ void gemv_awq_w4a16(nb::ndarray<> x, nb::ndarray<> qweight, nb::ndarray<> wscale check_hip_launch(); } -void svdquant_lora_down_impl(nb::ndarray<> x, nb::ndarray<> lora_down, nb::ndarray<> lora_act, - int M, int K, int R, uintptr_t stream_ptr, bool use_wmma) { - constexpr const char* kFn = "svdquant_lora_down"; +static void svdquant_lora_down_impl(const char* kFn, nb::ndarray<> x, + nb::ndarray<> lora_down, nb::ndarray<> lora_act, int M, int K, + int R, uintptr_t stream_ptr, bool use_wmma) { require_nonneg(M, kFn, "M"); require_nonneg(K, kFn, "K"); require_nonneg(R, kFn, "R"); @@ -979,12 +979,14 @@ void svdquant_lora_down_impl(nb::ndarray<> x, nb::ndarray<> lora_down, nb::ndarr 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) { - svdquant_lora_down_impl(x, lora_down, lora_act, M, K, R, stream_ptr, false); + svdquant_lora_down_impl("svdquant_lora_down", x, lora_down, lora_act, M, K, R, stream_ptr, + false); } void svdquant_lora_down_wmma(nb::ndarray<> x, nb::ndarray<> lora_down, nb::ndarray<> lora_act, int M, int K, int R, uintptr_t stream_ptr) { - svdquant_lora_down_impl(x, lora_down, lora_act, M, K, R, stream_ptr, true); + svdquant_lora_down_impl("svdquant_lora_down_wmma", x, lora_down, lora_act, M, K, R, + stream_ptr, true); } void svdquant_quantize(nb::ndarray<> x, nb::ndarray<> smooth, nb::ndarray<> q, diff --git a/tests/test_hip_wmma.py b/tests/test_hip_wmma.py index df1a9d32..8e93bc52 100644 --- a/tests/test_hip_wmma.py +++ b/tests/test_hip_wmma.py @@ -2015,13 +2015,46 @@ def test_svdquant_lora_down_scalar_fallback(hip, monkeypatch): assert torch.count_nonzero(lora_act[m:]) == 0 +@needs_wmma +def test_svdquant_lora_down_wmma_entry_falls_back_for_fp32(hip): + """Float32 operands take the scalar path inside the WMMA launcher.""" + torch.manual_seed(0) + m, k, rank = 19, 128, 17 + x = torch.randn(m, k, device=DEV, dtype=torch.float32) + smooth = torch.ones(k, device=DEV, dtype=torch.float32) + lora_down = torch.randn(k, rank, device=DEV, dtype=torch.float32) + + _, _, lora_act = hip.quantize_svdquant_w4a4( + x, smooth, lora_down, pad_size=16 + ) + ref = x @ lora_down + + torch.testing.assert_close(lora_act[:m], ref, rtol=2e-3, atol=2e-3) + assert torch.count_nonzero(lora_act[m:]) == 0 + + +def test_svdquant_lora_down_wmma_reports_entry_point(hip): + """Validation errors identify the WMMA binding rather than the scalar binding.""" + m, k, rank = 1, 64, 1 + x = torch.zeros(m, k, device=DEV, dtype=torch.float16) + lora_down = torch.zeros(k, rank, device=DEV, dtype=torch.float16) + bad_lora_act = torch.zeros(m, rank, device=DEV, dtype=torch.float16) + + with pytest.raises(RuntimeError, match="svdquant_lora_down_wmma: lora_act"): + hip._C.svdquant_lora_down_wmma( + hip._dl(x), hip._dl(lora_down), hip._dl(bad_lora_act), + m, k, rank, hip._stream(x), + ) + + +@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): +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, k = 3, 128 + m = 3 values = torch.arange(k, device=DEV).repeat(m, 1) if act_unsigned: values = values.remainder(16) From 71fd9edb197eda150fd288d6a45c12a6944b0b97 Mon Sep 17 00:00:00 2001 From: tangzzycc <3081129260@qq.com> Date: Sat, 22 Aug 2026 17:43:38 +0800 Subject: [PATCH 3/3] perf(hip): use torch.mm for SVDQuant LoRA-down Remove the custom scalar and WMMA LoRA-down kernels and bindings. Use the same backend matmul path as CUDA while retaining the optimized INT4 activation quantizer. --- comfy_kitchen/backends/hip/__init__.py | 15 +- .../backends/hip/dlpack_bindings.cpp | 46 ------ .../backends/hip/ops/svdquant_w4a4.hip | 140 ------------------ tests/test_hip_wmma.py | 68 +-------- 4 files changed, 15 insertions(+), 254 deletions(-) diff --git a/comfy_kitchen/backends/hip/__init__.py b/comfy_kitchen/backends/hip/__init__.py index f3ada51e..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,11 +1266,12 @@ def quantize_svdquant_w4a4( _dl(xc), _dl(smooth), _dl(q), _dl(ascales), m, m_pad, k, act_unsigned, _stream(x), ) - lora_down_op = _C.svdquant_lora_down_wmma if has_wmma() else _C.svdquant_lora_down - lora_down_op( - _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 ff90c281..efe272e4 100644 --- a/comfy_kitchen/backends/hip/dlpack_bindings.cpp +++ b/comfy_kitchen/backends/hip/dlpack_bindings.cpp @@ -83,10 +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_lora_down_wmma_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*, @@ -949,46 +945,6 @@ void gemv_awq_w4a16(nb::ndarray<> x, nb::ndarray<> qweight, nb::ndarray<> wscale check_hip_launch(); } -static void svdquant_lora_down_impl(const char* kFn, nb::ndarray<> x, - nb::ndarray<> lora_down, nb::ndarray<> lora_act, int M, int K, - int R, uintptr_t stream_ptr, bool use_wmma) { - 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"); - - const int x_code = map_dtype_to_code(x.dtype()); - const int d_code = map_dtype_to_code(lora_down.dtype()); - if (use_wmma) { - launch_svdquant_lora_down_wmma_kernel( - x.data(), lora_down.data(), lora_act.data(), M, K, R, x_code, d_code, - reinterpret_cast(stream_ptr)); - } else { - launch_svdquant_lora_down_kernel( - x.data(), lora_down.data(), lora_act.data(), M, K, R, x_code, d_code, - reinterpret_cast(stream_ptr)); - } - 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) { - svdquant_lora_down_impl("svdquant_lora_down", x, lora_down, lora_act, M, K, R, stream_ptr, - false); -} - -void svdquant_lora_down_wmma(nb::ndarray<> x, nb::ndarray<> lora_down, nb::ndarray<> lora_act, - int M, int K, int R, uintptr_t stream_ptr) { - svdquant_lora_down_impl("svdquant_lora_down_wmma", x, lora_down, lora_act, M, K, R, - stream_ptr, true); -} - 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) { @@ -1560,8 +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_lora_down_wmma", &svdquant_lora_down_wmma); 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 851ade36..7c27b76f 100644 --- a/comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip +++ b/comfy_kitchen/backends/hip/ops/svdquant_w4a4.hip @@ -26,112 +26,6 @@ namespace comfy::hip_backend { constexpr int kSvdGroup = 64; // quantization group, in elements constexpr int kSvdGroupBytes = 32; // ... and in packed int4 bytes -// Scalar fallback for FP32 and mixed input dtypes. -__global__ __launch_bounds__(256) void lora_down_scalar_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; -} - -// x[M, K] @ lora_down[K, R] -> lora_act[M, R], with fp32 accumulation. -// -// One wave computes a 16-row by (TN * 16)-column output tile. The A fragment is -// contiguous in x. lora_down is row-major KxR rather than the transposed RxK -// layout consumed by the WMMA instruction, so each lane gathers one output -// column across the K step. For a fixed K element the 16 fragment rows are -// adjacent in memory, keeping those gathers coalesced without an LDS transpose. -template -__global__ __launch_bounds__(kWave) void lora_down_wmma_kernel( - const typename Mma::Elem* __restrict__ x, - const typename Mma::Elem* __restrict__ lora_down, - float* __restrict__ lora_act, int M, int K, int R) { - - const int lane = threadIdx.x; - const int frag_r = frag_row(lane); - const int m0 = blockIdx.y * 16; - const int r0 = blockIdx.x * (TN * 16); - - typename Mma::Acc acc[TN]; -#pragma unroll - for (int n = 0; n < TN; ++n) { - acc[n] = Mma::zero(); - } - - for (int k0 = 0; k0 < K; k0 += 16) { - typename Mma::Frag a{}; - const int row = m0 + frag_r; - if (row < M) { - a = load_frag_16bit( - x + static_cast(row) * K + k0, lane); - } - -#pragma unroll - for (int n = 0; n < TN; ++n) { - typename Mma::Frag b{}; - const int col = r0 + n * 16 + frag_r; - if (col < R) { - const int kbase = k0 + Mma::frag_base(lane); -#pragma unroll - for (int i = 0; i < Mma::kFragElems; ++i) { - b[i] = lora_down[static_cast(kbase + i) * R + col]; - } - } - acc[n] = Mma::mma(a, b, acc[n]); - } - } - - const int col_lane = acc_col(lane); -#pragma unroll - for (int e = 0; e < 8; ++e) { - const int row = m0 + acc_row(lane, e); - if (row >= M) continue; -#pragma unroll - for (int n = 0; n < TN; ++n) { - const int col = r0 + n * 16 + col_lane; - if (col < R) { - lora_act[static_cast(row) * R + col] = Mma::get(acc[n], e); - } - } - } -} - -template -void launch_lora_down_wmma( - const void* x, const void* lora_down, float* lora_act, - int M, int K, int R, hipStream_t stream) { - - using Elem = typename Mma::Elem; - if (R <= 16) { - constexpr int TN = 1; - const dim3 grid((R + TN * 16 - 1) / (TN * 16), (M + 15) / 16); - lora_down_wmma_kernel<<>>( - static_cast(x), static_cast(lora_down), - lora_act, M, K, R); - } else if (R <= 32) { - constexpr int TN = 2; - const dim3 grid((R + TN * 16 - 1) / (TN * 16), (M + 15) / 16); - lora_down_wmma_kernel<<>>( - static_cast(x), static_cast(lora_down), - lora_act, M, K, R); - } else { - constexpr int TN = 4; - const dim3 grid((R + TN * 16 - 1) / (TN * 16), (M + 15) / 16); - lora_down_wmma_kernel<<>>( - static_cast(x), static_cast(lora_down), - lora_act, M, K, R); - } -} - // Scalar fallback for operands with different dtypes. template __global__ __launch_bounds__(256) void svdquant_quant_scalar_kernel( @@ -417,40 +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 || R == 0) { - return; // a zero-block launch is hipErrorInvalidConfiguration - } - lora_down_scalar_kernel<<>>( - x, lora_down, static_cast(lora_act), M, K, R, x_code, d_code); -} - -extern "C" void launch_svdquant_lora_down_wmma_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 || R == 0) { - return; // a zero-block launch is hipErrorInvalidConfiguration - } - if ((K & 15) == 0 && x_code == 1 && d_code == 1) { - launch_lora_down_wmma( - x, lora_down, static_cast(lora_act), M, K, R, stream); - } else if ((K & 15) == 0 && x_code == 2 && d_code == 2) { - launch_lora_down_wmma( - x, lora_down, static_cast(lora_act), M, K, R, stream); - } else { - launch_svdquant_lora_down_kernel( - x, lora_down, lora_act, M, K, R, x_code, d_code, stream); - } -} - 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) { diff --git a/tests/test_hip_wmma.py b/tests/test_hip_wmma.py index 8e93bc52..d3fd6089 100644 --- a/tests/test_hip_wmma.py +++ b/tests/test_hip_wmma.py @@ -1975,13 +1975,11 @@ def test_svdquant_validates_its_operands(hip): assert torch.isfinite(out).all() -@needs_wmma -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("rank", [1, 7, 16, 17, 32, 33, 64, 65, 96, 128, 192, 256]) -def test_svdquant_lora_down_matches_fp32_reference(hip, dtype, rank): - """The WMMA LoRA-down path handles partial M/R tiles and accumulates in fp32.""" - torch.manual_seed(rank) - m, k = 19, 128 +@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) @@ -1989,64 +1987,12 @@ def test_svdquant_lora_down_matches_fp32_reference(hip, dtype, rank): _, _, lora_act = hip.quantize_svdquant_w4a4( x, smooth, lora_down, pad_size=16 ) - ref = x.float() @ lora_down.float() - - torch.testing.assert_close( - lora_act[:m], ref, rtol=2e-3, atol=2e-3 - ) - assert torch.count_nonzero(lora_act[m:]) == 0 - - -def test_svdquant_lora_down_scalar_fallback(hip, monkeypatch): - """Devices without matrix cores retain the scalar LoRA-down path.""" - torch.manual_seed(0) - m, k, rank = 5, 64, 7 - x = torch.randn(m, k, device=DEV, dtype=torch.float16) - smooth = torch.ones(k, device=DEV, dtype=torch.float16) - lora_down = torch.randn(k, rank, device=DEV, dtype=torch.float16) - monkeypatch.setattr(hip, "has_wmma", lambda: False) - - _, _, lora_act = hip.quantize_svdquant_w4a4( - x, smooth, lora_down, pad_size=16 - ) - ref = x.float() @ lora_down.float() + ref = torch.mm(x, lora_down).float() - torch.testing.assert_close(lora_act[:m], ref, rtol=2e-3, atol=2e-3) + torch.testing.assert_close(lora_act[:m], ref) assert torch.count_nonzero(lora_act[m:]) == 0 -@needs_wmma -def test_svdquant_lora_down_wmma_entry_falls_back_for_fp32(hip): - """Float32 operands take the scalar path inside the WMMA launcher.""" - torch.manual_seed(0) - m, k, rank = 19, 128, 17 - x = torch.randn(m, k, device=DEV, dtype=torch.float32) - smooth = torch.ones(k, device=DEV, dtype=torch.float32) - lora_down = torch.randn(k, rank, device=DEV, dtype=torch.float32) - - _, _, lora_act = hip.quantize_svdquant_w4a4( - x, smooth, lora_down, pad_size=16 - ) - ref = x @ lora_down - - torch.testing.assert_close(lora_act[:m], ref, rtol=2e-3, atol=2e-3) - assert torch.count_nonzero(lora_act[m:]) == 0 - - -def test_svdquant_lora_down_wmma_reports_entry_point(hip): - """Validation errors identify the WMMA binding rather than the scalar binding.""" - m, k, rank = 1, 64, 1 - x = torch.zeros(m, k, device=DEV, dtype=torch.float16) - lora_down = torch.zeros(k, rank, device=DEV, dtype=torch.float16) - bad_lora_act = torch.zeros(m, rank, device=DEV, dtype=torch.float16) - - with pytest.raises(RuntimeError, match="svdquant_lora_down_wmma: lora_act"): - hip._C.svdquant_lora_down_wmma( - hip._dl(x), hip._dl(lora_down), hip._dl(bad_lora_act), - m, k, rank, hip._stream(x), - ) - - @pytest.mark.parametrize("k", [128, 1088]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @pytest.mark.parametrize("act_unsigned", [False, True])