hub / github.com/FlashSampling/FlashSampling / matmul_kernel

Function matmul_kernel

src/fused_mm_sampling/persistent_matmul.py:93–153 · view source on GitHub ↗

(
    a_ptr,
    b_ptr,
    c_ptr,  #
    M,
    N,
    K,  #
    stride_am,
    stride_ak,  #
    stride_bk,
    stride_bn,  #
    stride_cm,
    stride_cn,  #
    BLOCK_SIZE_M: tl.constexpr,  #
    BLOCK_SIZE_N: tl.constexpr,  #
    BLOCK_SIZE_K: tl.constexpr,  #
    GROUP_SIZE_M: tl.constexpr,  #
)

Source from the content-addressed store, hash-verified

91	)
92	@triton.jit(launch_metadata=_matmul_launch_metadata)
93	def matmul_kernel(
94	a_ptr,
95	b_ptr,
96	c_ptr, #
97	M,
98	N,
99	K, #
100	stride_am,
101	stride_ak, #
102	stride_bk,
103	stride_bn, #
104	stride_cm,
105	stride_cn, #
106	BLOCK_SIZE_M: tl.constexpr, #
107	BLOCK_SIZE_N: tl.constexpr, #
108	BLOCK_SIZE_K: tl.constexpr, #
109	GROUP_SIZE_M: tl.constexpr, #
110	):
111	pid = tl.program_id(axis=0)
112	num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
113	num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
114	num_pid_in_group = GROUP_SIZE_M * num_pid_n
115	group_id = pid // num_pid_in_group
116	first_pid_m = group_id * GROUP_SIZE_M
117	group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
118	pid_m = first_pid_m + (pid % group_size_m)
119	pid_n = (pid % num_pid_in_group) // group_size_m
120
121	start_m = pid_m * BLOCK_SIZE_M
122	start_n = pid_n * BLOCK_SIZE_N
123
124	offs_am = start_m + tl.arange(0, BLOCK_SIZE_M)
125	offs_bn = start_n + tl.arange(0, BLOCK_SIZE_N)
126	offs_am = tl.where(offs_am < M, offs_am, 0)
127	offs_bn = tl.where(offs_bn < N, offs_bn, 0)
128
129	offs_am = tl.max_contiguous(tl.multiple_of(offs_am, BLOCK_SIZE_M), BLOCK_SIZE_M)
130	offs_bn = tl.max_contiguous(tl.multiple_of(offs_bn, BLOCK_SIZE_N), BLOCK_SIZE_N)
131	offs_k = tl.arange(0, BLOCK_SIZE_K)
132	a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
133	b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
134
135	accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
136
137	for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
138	a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
139	b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
140	accumulator = tl.dot(a, b, accumulator)
141	a_ptrs += BLOCK_SIZE_K * stride_ak
142	b_ptrs += BLOCK_SIZE_K * stride_bk
143
144	if c_ptr.dtype.element_ty == tl.float8e4nv:
145	c = accumulator.to(tl.float8e4nv)
146	else:
147	c = accumulator.to(tl.float16)
148
149	offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
150	offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected