Add SVD quant for Quark

docs/prebuilt_kernels_guide.md

@@ -10,6 +10,7 @@
 | **RMSNorm** | `build_rmsnorm_module(N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | LDS-cached 3-pass pipeline |
 | **Softmax** | `build_softmax_module(M, N, dtype)` | Layout API (`@flyc.kernel`) | f32, f16, bf16 | Online softmax, adaptive block size |
 | **GEMM** | `compile_preshuffle_gemm(...)` | `@flyc.kernel` | fp8, int8, fp16, bf16 | Preshuffle B, ping-pong LDS, MFMA 16x16 |
+| **A8W4 SVDQuant GEMM** | `compile_preshuffle_gemm_a8w4(...)` | `@flyc.kernel` | MXFP8 × MXFP4 | gfx950 residual GEMM with optional fused low-rank up-projection |
 | **FlashAttention** | `build_flash_attn_func_module(...)` | `@flyc.kernel` | bf16, f16 (any arch); fp8 e4m3fn (gfx950, D=128, dense) | Dual-wave SWP fwd, GQA/MQA, causal, descale ABI |
 
 > **Note on API styles**: All kernels use the `@flyc.kernel`/`@flyc.jit` API from `flydsl.compiler` and `flydsl.expr` (`python/flydsl/`).
@@ -183,6 +184,50 @@ the separate `compile_mxfp4_gemm` in `kernels/gemm/gemm_fp8fp4_gfx1250.py` is th
 distinct gfx1250 kernel. `batch>1` runs a strided-batched GEMM over `grid.z`.
 Covered by `tests/kernels/test_preshuffle_gemm.py`.
 
+### 3.2 A8W4 SVDQuant Preshuffle GEMM (`kernels/gemm/preshuffle_gemm_a8w4.py`)
+
+This standalone gfx950 kernel consumes an MXFP8 activation and a preshuffled
+MXFP4 weight, both with per-1x32 E8M0 scales. It is intended for integrations
+that precompute a low-rank down-projection `d = x @ L1.T` outside FlyDSL and
+need the up-projection fused into the residual GEMM:
+
+```python
+from kernels.gemm.preshuffle_gemm_a8w4 import compile_preshuffle_gemm_a8w4
+
+launch_fn = compile_preshuffle_gemm_a8w4(
+    M=0,
+    N=3072,
+    K=3072,
+    tile_m=32,
+    tile_n=128,
+    tile_k=128,
+    out_dtype="bf16",
+    epilogue="svd_bias",
+    rank=32,
+)
+```
+
+`epilogue` supports `"none"`, the bias/activation variants, `"svd"`, and
+`"svd_bias"`. The SVD modes compute `C += d @ L2.T` in f32 before conversion
+to the output type; `"svd_bias"` additionally adds a per-N bias. `rank` is a
+compile-time value. When `rank` is divisible by 16, the default
+`svd_use_mfma=True` uses BF16 MFMA for the up-projection; other ranks use the
+scalar fallback. SVD epilogues do not support `use_cshuffle_epilog=True`.
+
+The launcher ABI is:
+
+```python
+launch_fn(C, A, B, scale_a, scale_b, bias, d, L2, M, N, stream)
+```
+
+For `"none"` and non-SVD epilogues, pass empty tensors for unused `bias`, `d`,
+and `L2`. `A` is row-major MXFP8, `B` is 16x16-preshuffled packed MXFP4, and
+both scales use the matching `shuffle_scale_w4` layout. `K` must be divisible
+by the selected `tile_k` and by 256; `N` must be divisible by `tile_n`. The
+kernel handles ragged `M` through its runtime buffer bounds. The W4A4 factory
+`compile_preshuffle_gemm_w4(...)` exposes the same epilogue and launcher
+contract.
+
 **Pipeline details:**
 - **lds_stage=2 (ping-pong)**: Two LDS buffers for A tiles. Cross-tile A0 prefetch overlaps VMEM with LDS reads
 - **lds_stage=1 (single)**: CK-style intrawave schedule with single LDS buffer

kernels/gemm/mxfp4_preshuffle.py

@@ -8,7 +8,8 @@
 import flydsl.compiler as flyc
 import flydsl.expr as fx
 from flydsl._mlir.dialects import fly
-from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr import buffer_ops, const_expr, gpu, range_constexpr, rocdl
+from flydsl.expr import math as fx_math
 from flydsl.expr.typing import (
     BFloat16,
     Constexpr,
@@ -61,6 +62,7 @@ def launch_gemm(
     arg_b: fx.Pointer,
     arg_scale_a: fx.Pointer,
     arg_scale_b: fx.Pointer,
+    arg_bias: fx.Pointer,
     i32_m: fx.Int32,
     i32_n: fx.Int32,
     stream: fx.Stream,
@@ -81,6 +83,7 @@ def launch_gemm(
     c_batch_stride: Constexpr[int],
     waves_per_eu: Constexpr[int],
     xcd_swizzle: Constexpr[int] = 0,
+    epilogue: Constexpr[str] = "none",
 ):
     """Direct @flyc.jit launcher. Operands are fx.Pointer (pass ptr_arg(t): raw data_ptr, no
     per-launch DLPack). Compile once with flyc.compile, then cf(*runtime). a_dtype fp4/fp6/fp8
@@ -95,6 +98,13 @@ def launch_gemm(
     else:
         out_elem = Float16
 
+    # Fused epilogue: per-N bias add (+ optional activation) folded into the C store.
+    # "none"/"bias"/"bias_relu"/"bias_silu"/"bias_gelu" (tanh-approx GELU).
+    _has_bias = epilogue in ("bias", "bias_relu", "bias_silu", "bias_gelu")
+    _has_relu = epilogue == "bias_relu"
+    _has_silu = epilogue == "bias_silu"
+    _has_gelu = epilogue == "bias_gelu"
+
     # Row sizes + read_a fragment layout (i32 units): fp6/fp8 read two b128 halves -> i32[A_NDW], fp4 one -> i32[4].
     if const_expr(a_dtype == "fp4"):  # 2 codes/byte
         a_row_bytes, A_ROW_B = K // 2, BK // 2
@@ -149,6 +159,7 @@ def kernel_gemm(
         arg_b: fx.Int64,
         arg_scale_a: fx.Int64,
         arg_scale_b: fx.Int64,
+        arg_bias: fx.Int64,
         i32_m: fx.Int32,
         i32_n: fx.Int32,
     ):
@@ -454,11 +465,41 @@ def hot_loop_scheduler():
         c_copy = fx.make_copy_atom(fx.rocdl.BufferCopy16b(), out_elem)
         c_rstride = fx.Int32(c_stride)
         col_w = by_n + wave * (BN // 4) + lane_mod_16
+        # Per-N bias resource (out_elem dtype, length N) for the fused epilogue.
+        bias_rsrc = None
+        if const_expr(_has_bias):
+            bias_ptr_ty = fx.PointerType.get(out_elem.ir_type, address_space=fx.AddressSpace.Global, alignment=2)
+            bias_view = fx.Tensor(fx.make_view(fx.inttoptr(bias_ptr_ty, arg_bias), fx.make_layout(N, 1)))
+            bias_rsrc = buffer_ops.create_buffer_resource(bias_view, max_size=True)
         for mi in range_constexpr(m_chunks):
             row_m = bx_m + mi * 16 + lane_div_16 * 4
             for ni in range_constexpr(num_acc_n):
                 col = col_w + ni * 16
-                acc = Vec(accs[mi * num_acc_n + ni]).to(out_elem)
+                if const_expr(_has_bias):
+                    # bias is per output column (N); apply + activation in f32.
+                    acc_f = Vec(accs[mi * num_acc_n + ni])  # f32 accumulator
+                    bval = fx.Float32(buffer_ops.buffer_load(bias_rsrc, col, vec_width=1, dtype=out_elem.ir_type))
+                    outs = []
+                    for ii in range_constexpr(4):
+                        v = fx.Float32(acc_f[ii]) + bval
+                        if const_expr(_has_relu):
+                            v = fx.Float32(v).maximumf(fx.Float32(0.0))
+                        elif const_expr(_has_silu):
+                            v = v * (fx.Float32(1.0) / (fx.Float32(1.0) + fx_math.exp(v * fx.Float32(-1.0))))
+                        elif const_expr(_has_gelu):
+                            # tanh-approx GELU, overflow-safe (non-positive exponent).
+                            x3 = v * v * v
+                            y = fx.Float32(0.7978845608) * (v + fx.Float32(0.044715) * x3)
+                            abs_y = fx.Float32(y).maximumf(fx.Float32(0.0) - y)
+                            e = fx_math.exp(fx.Float32(-2.0) * abs_y)
+                            denom = fx.Float32(1.0) + e
+                            one_plus_tanh = (y >= fx.Float32(0.0)).select(
+                                fx.Float32(2.0) / denom, (fx.Float32(2.0) * e) / denom)
+                            v = fx.Float32(0.5) * v * one_plus_tanh
+                        outs.append(v)
+                    acc = Vec.from_elements(outs, fx.Float32).to(out_elem)
+                else:
+                    acc = Vec(accs[mi * num_acc_n + ni]).to(out_elem)
                 for ii in range_constexpr(4):
                     cf = fx.make_rmem_tensor(1, out_elem)
                     cf.store(Vec.from_elements([acc[ii]], out_elem))
@@ -470,6 +511,7 @@ def hot_loop_scheduler():
     b_addr = fx.Int64(fx.ptrtoint(arg_b))
     sa_addr = fx.Int64(fx.ptrtoint(arg_scale_a))
     sb_addr = fx.Int64(fx.ptrtoint(arg_scale_b))
+    bias_addr = fx.Int64(fx.ptrtoint(arg_bias))
     if const_expr(waves_per_eu > 0):
         wpe = waves_per_eu
     else:
@@ -482,6 +524,7 @@ def hot_loop_scheduler():
         b_addr,
         sa_addr,
         sb_addr,
+        bias_addr,
         i32_m,
         i32_n,
         value_attrs={"rocdl.waves_per_eu": wpe},