HIP: optimize SVDQuant INT4 activation quantization

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
 
 

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<int64_t>(M) * K, kFn, "x");
-    require_len(lora_down, static_cast<int64_t>(K) * R, kFn, "lora_down");
-    require_len(lora_act, static_cast<int64_t>(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<hipStream_t>(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);
 }

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<int64_t>(m) * K + k, x_code) *
-               load_in(lora_down, static_cast<int64_t>(k) * R + r, d_code);
-    }
-    lora_act[static_cast<int64_t>(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 <bool UNSIGNED>
-__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 <typename Elem, bool UNSIGNED>
+__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<int64_t>(m) * K + kbase;
+
+        float values[kElemsPerThread];
+        float absmax = 0.0f;
+#pragma unroll
+        for (int i = 0; i < kElemsPerThread; ++i) {
+            values[i] = static_cast<float>(x[xbase + i]) /
+                        static_cast<float>(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<int64_t>(g) * M_pad + m] = static_cast<Elem>(scale);
+        }
+
+        int quantized[kElemsPerThread];
+#pragma unroll
+        for (int i = 0; i < kElemsPerThread; ++i) {
+            int v = static_cast<int>(rintf(values[i] * inv));
+            quantized[i] = v < kQMin ? kQMin : (v > kQMaxInt ? kQMaxInt : v);
+        }
+
+        const uint16_t packed =
+            static_cast<uint16_t>((quantized[0] & 0xF) |
+                                  ((quantized[1] & 0xF) << 4) |
+                                  ((quantized[2] & 0xF) << 8) |
+                                  ((quantized[3] & 0xF) << 12));
+        int8_t* qrow = q + static_cast<int64_t>(m) * (K / 2) + g * kSvdGroupBytes;
+        reinterpret_cast<uint16_t*>(qrow)[lane] = packed;
+    }
+}
+
+template <bool UNSIGNED>
+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<float, UNSIGNED><<<M, 256, 0, stream>>>(
+                static_cast<const float*>(x), static_cast<const float*>(smooth), q,
+                static_cast<float*>(ascales), M, M_pad, K);
+            return;
+        }
+        if (x_code == 1) {
+            svdquant_quant_kernel<__half, UNSIGNED><<<M, 256, 0, stream>>>(
+                static_cast<const __half*>(x), static_cast<const __half*>(smooth), q,
+                static_cast<__half*>(ascales), M, M_pad, K);
+            return;
+        }
+        if (x_code == 2) {
+            svdquant_quant_kernel<__bf16, UNSIGNED><<<M, 256, 0, stream>>>(
+                static_cast<const __bf16*>(x), static_cast<const __bf16*>(smooth), q,
+                static_cast<__bf16*>(ascales), M, M_pad, K);
+            return;
+        }
+    }
+
+    svdquant_quant_scalar_kernel<UNSIGNED><<<M, 256, 0, stream>>>(
+        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 <typename OutT, bool UNSIGNED>
 __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<<<M, 256, 0, stream>>>(x, lora_down, static_cast<float*>(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<true><<<M, 256, 0, stream>>>(
-            x, smooth, static_cast<int8_t*>(q), ascales, M, M_pad, K, x_code, s_code);
+        launch_svdquant_quant<true>(
+            x, smooth, static_cast<int8_t*>(q), ascales,
+            M, M_pad, K, x_code, s_code, stream);
     } else {
-        svdquant_quant_kernel<false><<<M, 256, 0, stream>>>(
-            x, smooth, static_cast<int8_t*>(q), ascales, M, M_pad, K, x_code, s_code);
+        launch_svdquant_quant<false>(
+            x, smooth, static_cast<int8_t*>(q), ascales,
+            M, M_pad, K, x_code, s_code, stream);
     }
 }
 

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)