[BENCH] 2% speedup from aligning matmul_ogs metadata (#6882)

This allows unmasked vectorised stores

Author: apgoucher

python/triton_kernels/tests/test_routing.py

@@ -52,11 +52,16 @@ def test_op(n_tokens, n_expts_tot, n_expts_act, block_m, device):
     ref_routing_data, ref_gather, ref_scatter = routing_torch(ref_logits, n_expts_act)
     tri_routing_data, tri_gather, tri_scatter = routing(tri_logits, n_expts_act)
     ref_metadata = ref_expt_data(ref_routing_data, n_tokens * n_expts_act, block_m)
-    tri_metadata = compute_metadata(tri_routing_data, n_tokens * n_expts_act, block_m).buffer
+    tri_metadata = compute_metadata(tri_routing_data, n_tokens * n_expts_act, block_m)
 
     assert_close(ref_routing_data.gate_scal, tri_routing_data.gate_scal, 2e-2, 4e-3)
     assert_equal(ref_routing_data.expt_hist, tri_routing_data.expt_hist)
-    assert_equal(ref_metadata, tri_metadata)
+
+    assert_equal(ref_metadata[:n_expts_tot], tri_metadata.hist)
+    assert_equal(ref_metadata[n_expts_tot:2 * n_expts_tot + 1], tri_metadata.offs)
+    assert_equal(ref_metadata[3 * n_expts_tot + 1], tri_metadata.offs_sum)
+    assert_equal(ref_metadata[3 * n_expts_tot + 2:], tri_metadata.blocks)
+
     assert ref_routing_data.n_expts_tot == ref_routing_data.n_expts_tot
     assert ref_routing_data.n_expts_act == ref_routing_data.n_expts_act
 

python/triton_kernels/triton_kernels/matmul_ogs.py

@@ -548,8 +548,9 @@ def matmul_ogs(x, w, bias,
         x, w, gather_indx, scatter_indx, routing_data, opt_flags, preprocessing_features
     )
     if expt_data.buffer is not None:
-        assert expt_data.buffer.shape[0] == 3*n_expts_tot + 2 + grid_m, \
-            f"invalid expt_data, {expt_data.buffer.shape}, {n_expts_tot=}, {grid_m=}"
+        assert expt_data.hist.shape[0] == n_expts_tot, "invalid expt_data"
+        assert expt_data.offs.shape[0] == n_expts_tot + 1, "invalid expt_data"
+        assert expt_data.blocks.shape[0] == grid_m, "invalid expt_data"
     # matrix multiplication
     n_cta = batch_size * grid_m * grid_n * opt_flags.split_k
     n_cta = min(target_info.num_sms(), n_cta) if opt_flags.is_persistent else n_cta

python/triton_kernels/triton_kernels/matmul_ogs_details/metadata.py

@@ -14,31 +14,49 @@ class ExptData:
 
 
 @triton.jit
-def _matmul_metadata_memset(Hist, n_expts_tot, MDHist, MDTokStarts, MDTileStarts, MDTileInfo, md_n_tiles,
-                            BLOCK: tl.constexpr, TILE_DIM: tl.constexpr):
+def _matmul_metadata_memset(Hist, n_expts_tot, MDTokStarts, MDTileStarts, MDTileInfo, BLOCK: tl.constexpr,
+                            TILE_DIM: tl.constexpr, extra_block: tl.constexpr):
     pid = tl.program_id(0)
+
+    TileInfoOut = MDTileInfo + (pid - 1) * BLOCK + tl.arange(0, BLOCK)
+
     # if pid == 0 - initialize cumsums
     if pid == 0:
         x_tok = tl.zeros([BLOCK], dtype=MDTokStarts.dtype.element_ty)
         x_tile = tl.zeros([BLOCK], dtype=MDTileStarts.dtype.element_ty)
-        tl.store(MDTokStarts, 0)
-        tl.store(MDTileStarts, 0)
+
+        Tok_ptrs = MDTokStarts + tl.arange(0, BLOCK)
+        Tile_ptrs = MDTileStarts + tl.arange(0, BLOCK)
+
         for i in range(0, n_expts_tot, BLOCK):
             offs_n = tl.arange(0, BLOCK) + i
-            mask = offs_n < n_expts_tot
-            hist_tok = tl.load(Hist + offs_n, mask=mask)
+            if extra_block:
+                # we need an extra block at the end just to contain the final
+                # sum; this only happens if our total number of experts is an
+                # exact multiple of BLOCK, obviating the need for any masking
+                hist_tok = tl.load(Hist + offs_n)
+            else:
+                mask = offs_n < n_expts_tot
+                hist_tok = tl.load(Hist + offs_n, mask=mask, other=0)
             hist_tile = tl.cdiv(hist_tok, TILE_DIM)
             tok_starts = tl.cumsum(hist_tok, 0) + x_tok
             x_tok += tl.sum(hist_tok, 0).to(MDTokStarts.dtype.element_ty)
             tile_starts = tl.cumsum(hist_tile, 0) + x_tile
             x_tile += tl.sum(hist_tile, 0).to(MDTileStarts.dtype.element_ty)
-            tl.store(MDHist + offs_n, hist_tok, mask=mask)
-            tl.store(MDTokStarts + 1 + offs_n, tok_starts, mask=mask)
-            tl.store(MDTileStarts + 1 + offs_n, tile_starts, mask=mask)
 
-    # initialize block data
-    offs = pid * BLOCK + tl.arange(0, BLOCK)
-    tl.store(MDTileInfo + offs, 0xffffffff, mask=offs < md_n_tiles)
+            tl.store(Tok_ptrs, tok_starts - hist_tok)
+            tl.store(Tile_ptrs, tile_starts - hist_tile)
+
+            Tok_ptrs += BLOCK
+            Tile_ptrs += BLOCK
+
+        if extra_block:
+            tl.store(Tok_ptrs, x_tok)
+            tl.store(Tile_ptrs, x_tile)
+
+    else:
+
+        tl.store(TileInfoOut, 0xffffffff)
 
 
 @triton.jit
@@ -60,7 +78,7 @@ def _matmul_metadata_compute(Hist, MDTileStarts, MDTileInfo, BLOCK: tl.constexpr
 def compute_metadata(routing_data, n_rows, block_m):
     if routing_data.expt_hist is None:
         return ExptData(None, None, None, None, None)
-    MEMSET_BLOCK = 512
+    MEMSET_BLOCK = 128
     HIST2_BLOCK_M = 512
     device = routing_data.expt_hist.device
     n_expts_tot = routing_data.n_expts_tot
@@ -69,21 +87,29 @@ def compute_metadata(routing_data, n_rows, block_m):
         grid_m = n_rows
     else:
         grid_m = n_expts_tot - 1 - ((n_expts_tot - n_rows - 1) // block_m)
-    metadata_size = 3 * n_expts_tot + 2 + grid_m
+
+    n_expts_pad = cdiv(n_expts_tot, MEMSET_BLOCK) * MEMSET_BLOCK
+    pad2 = cdiv(n_expts_tot + 1, MEMSET_BLOCK) * MEMSET_BLOCK
+    extra_block = (n_expts_pad != pad2)
+    pids = cdiv(grid_m, MEMSET_BLOCK) + 1
+
+    metadata_size = n_expts_pad + 2 * pad2 + MEMSET_BLOCK * (pids - 1)
+
     metadata = torch.empty(metadata_size, dtype=torch.int32, device=device)
-    md_hist = metadata[:n_expts_tot]
-    md_offs = metadata[n_expts_tot:n_expts_tot * 2 + 1]
-    md_offs_sum = metadata[3 * n_expts_tot + 2 - 1]
-    md_tile_starts = metadata[n_expts_tot * 2 + 1:n_expts_tot * 3 + 2]
-    md_tile_infos = metadata[n_expts_tot * 3 + 2:]
-    _matmul_metadata_memset[(cdiv(metadata_size, MEMSET_BLOCK), )](
-        routing_data.expt_hist, n_expts_tot, md_hist, md_offs, md_tile_starts, md_tile_infos, md_tile_infos.shape[0],
+
+    md_hist = routing_data.expt_hist[:n_expts_tot]
+    md_offs = metadata[:n_expts_tot + 1]
+    md_tile_starts = metadata[pad2:][:n_expts_tot + 1]
+    md_offs_sum = md_tile_starts[-1]
+    md_tile_infos = metadata[2 * pad2:][:grid_m]
+    _matmul_metadata_memset[(pids, )](
+        routing_data.expt_hist, n_expts_tot, md_offs, md_tile_starts, md_tile_infos,
         BLOCK=MEMSET_BLOCK,  # optimization parameters
         TILE_DIM=block_m,  # constants
-    )
+        extra_block=extra_block, num_warps=1)
     _matmul_metadata_compute[(n_expts_tot, )](
         routing_data.expt_hist, md_tile_starts, md_tile_infos,  # outputs
         BLOCK=HIST2_BLOCK_M,  # optimization parameters
         TILE_DIM=block_m,  # constants
-    )
+        num_warps=4)
     return ExptData(md_hist, md_offs, md_offs_sum, md_tile_infos, metadata)