[Gemm] Permute bias for concat_layout inside kernel if fp32

Author: tridao

quack/epi_ops.py

@@ -283,7 +283,7 @@ def begin(self, gemm, param, smem_tensor, ctx):
             tVgV = thr_copy.partition_S(gVec)
             tVsV = thr_copy.partition_D(smem_tensor)
             tVcV = thr_copy.partition_S(cute.make_identity_tensor(tile_dim))
-            limit = min(mVec.shape[0] - coord_idx * tile_dim, tile_dim)
+            limit = min(cute.size(mVec, mode=[0]) - coord_idx * tile_dim, tile_dim)
             pred = cute.make_rmem_tensor((1, cute.size(tVsV.shape[1])), Boolean)
             for m in cutlass.range(cute.size(tVsV.shape[1]), unroll_full=True):
                 pred[0, m] = tVcV[0, m] < limit

quack/gemm_act.py

@@ -37,6 +37,7 @@
     compile_gemm_kernel,
 )
 from quack.cache_utils import jit_cache
+import quack.layout_utils as layout_utils
 from quack.layout_utils import permute_gated_Cregs_b16
 from quack.activation import act_fn_map, gate_fn_map
 from quack.rounding import RoundingMode
@@ -67,6 +68,9 @@ def epi_to_underlying_arguments(self, args: EpilogueArguments, *, loc=None, ip=N
         self.cta_tile_shape_postact_mn = self.cta_tile_shape_mnk[:2]
         d = self._epi_ops_to_params_dict(args)
         d["act_fn"] = args.act_fn
+        for key in ("mRowVecBroadcast", "mColVecBroadcast"):
+            if key in self.concat_layout and key in d and d[key] is not None:
+                d[key] = layout_utils.concat_to_interleave(d[key], 1)
         return self.EpilogueParams(**d)
 
     # epi_get_tma_atoms, epi_smem_bytes_per_stage, epi_get_smem_struct,
@@ -210,6 +214,9 @@ def epi_to_underlying_arguments(
         )
         d = self._epi_ops_to_params_dict(args)
         d["act_fn"] = args.act_fn
+        for key in ("mRowVecBroadcast", "mColVecBroadcast"):
+            if key in self.concat_layout and key in d and d[key] is not None:
+                d[key] = layout_utils.concat_to_interleave(d[key], 1)
         return self.EpilogueParams(**d)
 
     @cute.jit

quack/gemm_default_epi.py

@@ -12,6 +12,7 @@
 from quack.gemm_sm100 import GemmSm100
 from quack.gemm_sm120 import GemmSm120
 from quack.rounding import RoundingMode
+import quack.layout_utils as layout_utils
 import quack.utils as utils
 
 
@@ -39,6 +40,9 @@ class EpilogueArguments(NamedTuple):
     def epi_to_underlying_arguments(self, args, *, loc=None, ip=None):
         self.rounding_mode = args.rounding_mode
         d = self._epi_ops_to_params_dict(args)
+        for key in ("mRowVecBroadcast", "mColVecBroadcast"):
+            if key in self.concat_layout and key in d and d[key] is not None:
+                d[key] = layout_utils.concat_to_interleave(d[key], 1)
         return self.EpilogueParams(**d)
 
     @cute.jit

quack/gemm_interface.py

@@ -60,11 +60,13 @@ def _concat_interleave(t):
     dim = -2 if t.stride(-1) == 1 else -1
     return t.unflatten(dim, (2, t.shape[dim] // 2)).transpose(dim - 1, dim).flatten(dim - 1, dim)
 
+
 def _concat_interleave_bias(t):
     """Interleave [gate; up] along last dim for bias vectors."""
     half = t.shape[-1] // 2
     return t.unflatten(-1, (2, half)).transpose(-2, -1).flatten(-2, -1)
 
+
 def default_config(device):
     cap = get_device_capacity(device)[0]
     if cap in [10, 11]:
@@ -204,16 +206,24 @@ def gemm_tuned(
         if dynamic_scheduler and get_device_capacity(A.device)[0] == 9
         else None
     )
+    # Handle bias concat layout: transform "bias" key to kernel-level key or permute data.
+    if concat_layout and "bias" in concat_layout:
+        if bias is not None and bias.dtype.itemsize >= 4:
+            # fp32: kernel permutes via layout; replace "bias" with the kernel-level key
+            concat_layout = tuple("mRowVecBroadcast" if k == "bias" else k for k in concat_layout)
+        else:
+            # No bias or sub-fp32: strip "bias" from concat_layout; permute data if needed
+            concat_layout = tuple(k for k in concat_layout if k != "bias")
+            if bias is not None:
+                bias = _concat_interleave_bias(bias)
     # When swap_ab, A↔B (out/C stay, but .mT flips their strides so the kernel
     # auto-detects the correct non-contiguous dim).
-    _swap_map = {"A": "B", "B": "A", "out": "out", "C": "C"}
+    _swap_map = {"A": "B", "B": "A", "out": "out", "C": "C", "mRowVecBroadcast": "mColVecBroadcast"}
     swapped_concat = (
         tuple(_swap_map.get(k, k) for k in concat_layout)
         if config.swap_ab and concat_layout
         else concat_layout
     )
-    if concat_layout and "B" in concat_layout and bias is not None:
-        bias = _concat_interleave_bias(bias)
     gemm_dispatch(
         A if not config.swap_ab else B,
         B if not config.swap_ab else A,
@@ -513,8 +523,8 @@ def gemm_ref(
             A = _concat_interleave(A)
         if "B" in concat_layout:
             B = _concat_interleave(B)
-            if bias is not None:
-                bias = _concat_interleave_bias(bias)
+        if "bias" in concat_layout and bias is not None:
+            bias = _concat_interleave_bias(bias)
     if cu_seqlens_m is None and cu_seqlens_k is None:
         fn = torch.bmm if A.ndim == 3 else torch.mm
         out = fn(A, B, out_dtype=out_dtype, out=out)
@@ -710,8 +720,8 @@ def gemm_add_ref(
             A = _concat_interleave(A)
         if "B" in concat_layout:
             B = _concat_interleave(B)
-            if bias is not None:
-                bias = _concat_interleave_bias(bias)
+        if "bias" in concat_layout and bias is not None:
+            bias = _concat_interleave_bias(bias)
         if "C" in concat_layout:
             C = _concat_interleave(C)
     if cu_seqlens_m is None and cu_seqlens_k is None:
@@ -1270,8 +1280,14 @@ def gemm_gated_tuned(
         PostAct = postact_out
     if bias is not None and bias.ndim == 1:
         bias = bias.unsqueeze(0)  # (L, N)
-    if concat_layout and "B" in concat_layout and bias is not None:
-        bias = _concat_interleave_bias(bias)
+    if concat_layout and "bias" in concat_layout:
+        if bias is not None and bias.dtype.itemsize >= 4:
+            bias_key = "mColVecBroadcast" if config.swap_ab else "mRowVecBroadcast"
+            concat_layout = tuple(bias_key if k == "bias" else k for k in concat_layout)
+        else:
+            concat_layout = tuple(k for k in concat_layout if k != "bias")
+            if bias is not None:
+                bias = _concat_interleave_bias(bias)
     dynamic_scheduler = dynamic_scheduler or config.is_dynamic_persistent
     tile_count_semaphore = (
         torch.zeros(1, dtype=torch.int32, device=A.device)

quack/linear.py

@@ -110,7 +110,7 @@ class _LinearGatedUntunedOps:
 
 
 class _LinearGatedConcatOps(_LinearGatedOps):
-    matmul_fwd_fn = partial(gemm_gated, concat_layout=("B",))
+    matmul_fwd_fn = partial(gemm_gated, concat_layout=("B", "bias"))
     matmul_bwd_dx = partial(gemm, dynamic_scheduler=True, concat_layout=("B",))
     matmul_bwd_dw = partial(gemm, dynamic_scheduler=True, concat_layout=("out",))
     matmul_bwd_dw_inplace = partial(
@@ -119,7 +119,7 @@ class _LinearGatedConcatOps(_LinearGatedOps):
 
 
 class _LinearGatedConcatUntunedOps(_LinearGatedUntunedOps):
-    matmul_fwd_fn = partial(gemm_gated, tuned=False, concat_layout=("B",))
+    matmul_fwd_fn = partial(gemm_gated, tuned=False, concat_layout=("B", "bias"))
     matmul_bwd_dx = partial(gemm, dynamic_scheduler=True, tuned=False, concat_layout=("B",))
     matmul_bwd_dw = partial(gemm, dynamic_scheduler=True, tuned=False, concat_layout=("out",))
     matmul_bwd_dw_inplace = partial(
@@ -274,10 +274,11 @@ def linear_gated_func(
 
 class DActLinearFunc(torch.autograd.Function):
     @staticmethod
-    def forward(ctx, preact, weight, x, activation, fuse_grad_accum, ops):
+    def forward(ctx, preact, weight, x, activation, bias, fuse_grad_accum, ops):
         """
         x: (..., in_features)
         weight: (out_features, in_features)
+        bias: (out_features,) or None
         out: (..., out_features)
         Takes in an extra preact argument which is the pre-activation, to be used in the backward pass.
         """
@@ -289,7 +290,7 @@ def forward(ctx, preact, weight, x, activation, fuse_grad_accum, ops):
             weight_og = weight
             batch_shape = x.shape[:-1]
             x = x.reshape(-1, x.shape[-1])
-            out = ops.matmul_fwd_fn(x, weight.T)
+            out = ops.matmul_fwd_fn(x, weight.T, bias=bias)
             # Store preact instead of x, we will recompute x (postact) in backward.
             # dpreact needs gemm_dact(dout, weight, preact) → needs both weight and preact.
             # dweight needs postact: if dpreact is also needed, postact comes from gemm_dact;
@@ -300,6 +301,8 @@ def forward(ctx, preact, weight, x, activation, fuse_grad_accum, ops):
                 ctx, preact, weight, weight_og, needs_x_w_grad=(need_weight, need_preact)
             )
             ctx.activation = activation
+            ctx.bias_dtype = bias.dtype if bias is not None else None
+            ctx.compute_dbias = bias is not None and ctx.needs_input_grad[4]
             return out.reshape(*batch_shape, out.shape[-1])
 
     @staticmethod
@@ -313,6 +316,7 @@ def backward(ctx, dout):
             preact, weight, weight_og = ctx.saved_tensors
             batch_shape = dout.shape[:-1]
             dout = _ensure_contiguous(dout.reshape(-1, dout.shape[-1]))
+            dbias = dout.sum(0, dtype=ctx.bias_dtype) if ctx.compute_dbias else None
             if ctx.needs_input_grad[0]:
                 # Need dpreact: gemm_dact(dout, weight, preact) → (dpreact, postact)
                 preact = preact.reshape(-1, preact.shape[-1])
@@ -331,17 +335,17 @@ def backward(ctx, dout):
             dweight = linear_bwd_compute_weight_grad(
                 ctx, dout, x, weight_og, ops.matmul_bwd_dw, ops.matmul_bwd_dw_inplace
             )
-            return dpreact, dweight, None, None, None, None
+            return dpreact, dweight, None, None, dbias, None, None
 
 
-def act_linear_func(preact, weight, x, activation, fuse_grad_accum=False, tuned=True):
+def act_linear_func(preact, weight, x, activation, bias=None, fuse_grad_accum=False, tuned=True):
     ops = _DActLinearOps if tuned else _DActLinearUntunedOps
-    return DActLinearFunc.apply(preact, weight, x, activation, fuse_grad_accum, ops)
+    return DActLinearFunc.apply(preact, weight, x, activation, bias, fuse_grad_accum, ops)
 
 
-def gated_linear_func(preact, weight, x, activation, fuse_grad_accum=False, tuned=True):
+def gated_linear_func(preact, weight, x, activation, bias=None, fuse_grad_accum=False, tuned=True):
     ops = _DGatedLinearOps if tuned else _DGatedLinearUntunedOps
-    return DActLinearFunc.apply(preact, weight, x, activation, fuse_grad_accum, ops)
+    return DActLinearFunc.apply(preact, weight, x, activation, bias, fuse_grad_accum, ops)
 
 
 class Linear(nn.Linear):

quack/mlp.py

@@ -207,6 +207,8 @@ def mlp_func(
     weight1,
     weight2,
     activation: str,
+    bias1=None,
+    bias2=None,
     fuse_grad_accum=False,
     tuned=True,
     recompute=False,
@@ -229,6 +231,7 @@ def mlp_func(
         x,
         weight1,
         activation,
+        bias=bias1,
         store_preact=torch.is_grad_enabled(),
         fuse_grad_accum=fuse_grad_accum,
         tuned=tuned,
@@ -239,6 +242,7 @@ def mlp_func(
         weight2,
         postact,
         activation=activation,
+        bias=bias2,
         fuse_grad_accum=fuse_grad_accum,
         tuned=tuned,
     )
@@ -292,9 +296,10 @@ def __init__(
         self.concat_layout = concat_layout
 
     def forward(self, input: Tensor) -> Tensor:
+        # Allow bias in the fused path during inference (fwd-only, no bwd).
+        bias_ok = not torch.is_grad_enabled() or (self.fc1.bias is None and self.fc2.bias is None)
         if (
-            self.fc1.bias is None
-            and self.fc2.bias is None
+            bias_ok
             and input.is_cuda
             and input.stride(-1) == 1
             and self.fc1.in_features % 8 == 0
@@ -306,6 +311,8 @@ def forward(self, input: Tensor) -> Tensor:
                 self.fc1.weight,
                 self.fc2.weight,
                 activation=self.activation,
+                bias1=self.fc1.bias,
+                bias2=self.fc2.bias,
                 fuse_grad_accum=self.fuse_grad_accum,
                 tuned=self.tuned,
                 recompute=self.recompute,

tests/test_linear.py

@@ -109,18 +109,23 @@ def test_gemm(m, k, n, input_dtype, A_major, B_major, out_major, swap_ab):
 @pytest.mark.parametrize("B_major", ["k", "n"])
 @pytest.mark.parametrize("A_major", ["k", "m"])
 @pytest.mark.parametrize("batched", [False, True])
+@pytest.mark.parametrize("bias_dtype", [None, torch.float32, torch.bfloat16])
 @pytest.mark.parametrize("concat_tensor", ["A", "B", "out"])
 @pytest.mark.parametrize("input_dtype", [torch.bfloat16])
 @pytest.mark.parametrize("n", [1504, 2048])
 @pytest.mark.parametrize("k", [736, 1024])
 @pytest.mark.parametrize("m", [960])
 def test_gemm_concat_layout(
-    m, k, n, input_dtype, concat_tensor, batched, A_major, B_major, out_major, swap_ab
+    m, k, n, input_dtype, concat_tensor, bias_dtype, batched, A_major, B_major, out_major, swap_ab
 ):
     """Test concat_layout on each tensor with all major/stride combinations."""
+    if bias_dtype is not None and batched:
+        pytest.skip("batched + bias not supported")
     device = "cuda"
     torch.random.manual_seed(0)
     concat = (concat_tensor,)
+    if bias_dtype is not None and concat_tensor == "B":
+        concat = ("B", "bias")
     batch_shape = (3,) if batched else ()
     A = torch.randn((*batch_shape, m, k), device=device, dtype=input_dtype) / math.sqrt(k)
     if A_major == "m":
@@ -131,20 +136,21 @@ def test_gemm_concat_layout(
     out = torch.empty((*batch_shape, m, n), device=device, dtype=input_dtype)
     if out_major == "m":
         out = out.mT.contiguous().mT
+    bias = torch.randn(n, device=device, dtype=bias_dtype) if bias_dtype is not None else None
     config = replace(default_config(torch.device(device)), swap_ab=swap_ab)
-    gemm_tuned.fn(A, B, out, config=config, concat_layout=concat)
+    gemm_tuned.fn(A, B, out, config=config, bias=bias, concat_layout=concat)
     # For ref: gemm_ref always produces n-major output and interleaves rows for concat="out".
     # When the kernel's out is m-major, the kernel interleaves columns instead.
     # Match by creating the ref with matching out major.
     if concat_tensor == "out" and out_major == "m":
         # Kernel interleaves columns. Ref: compute flat, then interleave columns.
-        out_ref_flat = gemm_ref(A.float(), B.float())
+        out_ref_flat = gemm_ref(A.float(), B.float(), bias=bias)
         out_ref = torch.cat([out_ref_flat[..., ::2], out_ref_flat[..., 1::2]], dim=-1)
-        out_pt_flat = gemm_ref(A, B)
+        out_pt_flat = gemm_ref(A, B, bias=bias)
         out_pt = torch.cat([out_pt_flat[..., ::2], out_pt_flat[..., 1::2]], dim=-1)
     else:
-        out_ref = gemm_ref(A.float(), B.float(), concat_layout=concat)
-        out_pt = gemm_ref(A, B, concat_layout=concat)
+        out_ref = gemm_ref(A.float(), B.float(), bias=bias, concat_layout=concat)
+        out_pt = gemm_ref(A, B, bias=bias, concat_layout=concat)
     assert (out - out_ref).abs().max() < 2 * (out_pt - out_ref).abs().max() + 1e-5
 
 
@@ -362,12 +368,12 @@ def test_gemm_add_inplace_alpha_beta(
 @pytest.mark.parametrize("store_preact", [True, False])
 @pytest.mark.parametrize("activation", ["swiglu", "swiglu_oai", "reglu", "geglu", "glu"])
 @pytest.mark.parametrize("input_dtype", [torch.bfloat16])
-@pytest.mark.parametrize("has_bias", [False, True])
+@pytest.mark.parametrize("bias_dtype", [None, torch.float32, torch.bfloat16])
 @pytest.mark.parametrize("out_features", [1504, 2048])
 @pytest.mark.parametrize("in_features", [736, 4096])
 @pytest.mark.parametrize("is_concat_layout_B", [False, True])
 def test_gemm_gated(
-    in_features, out_features, has_bias, input_dtype, activation, store_preact, is_concat_layout_B
+    in_features, out_features, bias_dtype, input_dtype, activation, store_preact, is_concat_layout_B
 ):
     """Test GEMM with gated activation forward computation."""
     device = "cuda"
@@ -381,8 +387,18 @@ def test_gemm_gated(
         / math.sqrt(in_features)
     ).requires_grad_()
     B = w.T
-    concat = ("B",) if is_concat_layout_B else None
-    bias = torch.randn(2 * out_features, device=device) if has_bias else None
+    bias = (
+        torch.randn(2 * out_features, device=device, dtype=bias_dtype)
+        if bias_dtype is not None
+        else None
+    )
+    concat = (
+        ("B", "bias")
+        if is_concat_layout_B and bias_dtype is not None
+        else ("B",)
+        if is_concat_layout_B
+        else None
+    )
     preact, postact = gemm_gated(
         x,
         B,