diff --git a/batchgen/attention/fused_kernels/qkv_wgmma.py b/batchgen/attention/fused_kernels/qkv_wgmma.py index e7673ce51..4f4cf9fc6 100644 --- a/batchgen/attention/fused_kernels/qkv_wgmma.py +++ b/batchgen/attention/fused_kernels/qkv_wgmma.py @@ -22,6 +22,16 @@ _module = None _qkv_wgmma_available: Optional[bool] = None +_arch: Optional[str] = None + + +def _get_arch() -> str: + """Cached device arch ('sm90a' / 'sm100').""" + global _arch + if _arch is None: + import batchgen_kernels + _arch = batchgen_kernels.get_device_arch() + return _arch def _check_wgmma_support() -> bool: @@ -66,6 +76,12 @@ def is_qkv_wgmma_available() -> bool: _qkv_wgmma_available = False return False + # On SM100 (Blackwell) the WGMMA .cu is not built; the fused QKV path is + # provided by the pure-Triton kernel (qkv_proj_rope), always available. + if _get_arch() == "sm100": + _qkv_wgmma_available = True + return True + _qkv_wgmma_available = _get_module() is not None return _qkv_wgmma_available @@ -102,6 +118,16 @@ def cuda_qkv_wgmma( (Q, K, V) as separate contiguous [M, *] BF16 tensors. When rope_cos/sin provided, Q and K have RoPE applied; V is unchanged. """ + # SM100 (Blackwell): WGMMA .cu is not built — use the pure-Triton port. + if _get_arch() == "sm100": + from batchgen_kernels.triton.qkv_proj_rope import qkv_proj_rope + _bias = bias if (bias is not None and bias.numel() > 0) else None + _cos = rope_cos if (rope_cos is not None and rope_cos.numel() > 0) else None + _sin = rope_sin if (rope_sin is not None and rope_sin.numel() > 0) else None + return qkv_proj_rope( + input, weight, _bias, q_size, kv_size, head_dim, _cos, _sin, + ) + mod = _get_module() assert mod is not None, "QKV WGMMA module not available" diff --git a/batchgen_kernels/triton/__init__.py b/batchgen_kernels/triton/__init__.py index 27874ab1f..63a1adf6d 100644 --- a/batchgen_kernels/triton/__init__.py +++ b/batchgen_kernels/triton/__init__.py @@ -33,3 +33,8 @@ from batchgen_kernels.triton.fused_dequant_gemm import fused_fp8_bf16_gemm from batchgen_kernels.triton.fused_q_absorb import fused_q_absorb_query_states from batchgen_kernels.triton.fused_out_absorb import fused_out_absorb_reshape +from batchgen_kernels.triton.qkv_proj_rope import ( + qkv_proj_split, + qkv_proj_rope, + apply_rope_qk, +) diff --git a/batchgen_kernels/triton/qkv_proj_rope.py b/batchgen_kernels/triton/qkv_proj_rope.py new file mode 100644 index 000000000..73da14563 --- /dev/null +++ b/batchgen_kernels/triton/qkv_proj_rope.py @@ -0,0 +1,243 @@ +"""QKV projection + split (+ optional RoPE) Triton kernel for SM100 (Blackwell). + +SM100 port of the SM90a WGMMA kernel `batchgen_kernels/src/attention/qkv_wgmma.cu`. +That kernel uses Hopper-only `wgmma.*` + warpgroup TMA, which do not exist on +sm_100a, so on Blackwell we replace it with a pure-Triton GEMM that Triton JIT-compiles +to UMMA. No `torch.matmul` / `F.linear` is used. + +Operation (decode/prefill): + + x[M, K] @ W_qkv[N, K]^T -> [M, N], N = q_size + kv_size + kv_size + +split column-wise into: + + Q[M, q_size] (cols [0, q_size)) + K[M, kv_size] (cols [q_size, q_size + kv_size)) + V[M, kv_size] (cols [q_size + kv_size, N)) + +RoPE (optional) is applied to Q and K only, never V (see `apply_rope_qk`). + +Design note (see batchgen-agent-metadata/.../incremental_compilation_contract.md and +blackwell-kernel-port-v1.md sub-task 5): RoPE is applied in a *separate* validated pass +(`apply_rope_qk`) rather than inline in the GEMM epilogue. At decode M the GEMM is +memory-bound on the [N, K] weight read, so the extra Q/K reread is negligible, and the +two-pass structure keeps each step independently testable — matching how BatchGen already +structures RoPE as dedicated kernels. +""" + +import torch +import triton +import triton.language as tl + + +# ────────────────────────────────────────────────────────────────────────────── +# Pass 1 — GEMM + column split into Q / K / V +# ────────────────────────────────────────────────────────────────────────────── + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_M': 32, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_stages=3, num_warps=4), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 64}, num_stages=3, num_warps=4), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=3, num_warps=8), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=2, num_warps=8), + triton.Config({'BLOCK_M': 16, 'BLOCK_N': 128, 'BLOCK_K': 128}, num_stages=3, num_warps=4), + ], + key=['M', 'K', 'N'], +) +@triton.jit +def _qkv_proj_split_kernel( + x_ptr, w_ptr, bias_ptr, + q_ptr, k_ptr, v_ptr, + M, K, N, q_size, kv_size, + stride_xm, stride_xk, + stride_wn, stride_wk, # w is [N, K] row-major: stride_wn=K, stride_wk=1 + stride_qm, stride_qn, + stride_km, stride_kn, + stride_vm, stride_vn, + HAS_BIAS: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + + m_mask = offs_m < M + n_mask = offs_n < N + + # x[BLOCK_M, K] @ w[N, K]^T -> acc[BLOCK_M, BLOCK_N], accumulate in fp32 + x_ptrs = x_ptr + (offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk) + # w is [N, K]; we need w[n, k] for the n-tile, contracted over k -> shape [BLOCK_K, BLOCK_N] + w_ptrs = w_ptr + (offs_n[None, :] * stride_wn + offs_k[:, None] * stride_wk) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k0 in range(0, K, BLOCK_K): + k_mask = (k0 + offs_k) < K + x_tile = tl.load(x_ptrs, mask=m_mask[:, None] & k_mask[None, :], other=0.0) + w_tile = tl.load(w_ptrs, mask=k_mask[:, None] & n_mask[None, :], other=0.0) + acc += tl.dot(x_tile, w_tile, out_dtype=tl.float32) + x_ptrs += BLOCK_K * stride_xk + w_ptrs += BLOCK_K * stride_wk + + if HAS_BIAS: + bias = tl.load(bias_ptr + offs_n, mask=n_mask, other=0.0).to(tl.float32) + acc += bias[None, :] + + out = acc.to(q_ptr.dtype.element_ty) + + kv_end = q_size + kv_size + + # Region masks (per-column). Robust to any BLOCK_N / boundary alignment. + is_q = offs_n < q_size + is_k = (offs_n >= q_size) & (offs_n < kv_end) + is_v = offs_n >= kv_end + + # Q store: local col = offs_n + q_cols = offs_n + q_ptrs = q_ptr + (offs_m[:, None] * stride_qm + q_cols[None, :] * stride_qn) + tl.store(q_ptrs, out, mask=m_mask[:, None] & (is_q & n_mask)[None, :]) + + # K store: local col = offs_n - q_size + k_cols = offs_n - q_size + k_ptrs = k_ptr + (offs_m[:, None] * stride_km + k_cols[None, :] * stride_kn) + tl.store(k_ptrs, out, mask=m_mask[:, None] & (is_k & n_mask)[None, :]) + + # V store: local col = offs_n - kv_end + v_cols = offs_n - kv_end + v_ptrs = v_ptr + (offs_m[:, None] * stride_vm + v_cols[None, :] * stride_vn) + tl.store(v_ptrs, out, mask=m_mask[:, None] & (is_v & n_mask)[None, :]) + + +def qkv_proj_split( + x: torch.Tensor, # [M, K] BF16 + w_qkv: torch.Tensor, # [N, K] BF16, N = q_size + 2*kv_size + bias: torch.Tensor | None, + q_size: int, + kv_size: int, +): + """Fused QKV GEMM + column split. Returns (Q, K, V) BF16, no RoPE.""" + assert x.dim() == 2 and w_qkv.dim() == 2 + M, K = x.shape + N, Kw = w_qkv.shape + assert Kw == K, f"K mismatch: x K={K}, w K={Kw}" + assert N == q_size + 2 * kv_size, f"N={N} != q_size({q_size}) + 2*kv_size({kv_size})" + + Q = torch.empty((M, q_size), dtype=x.dtype, device=x.device) + Kt = torch.empty((M, kv_size), dtype=x.dtype, device=x.device) + V = torch.empty((M, kv_size), dtype=x.dtype, device=x.device) + + HAS_BIAS = bias is not None + bias_arg = bias if HAS_BIAS else x # dummy ptr when unused + + grid = lambda meta: (triton.cdiv(M, meta['BLOCK_M']), triton.cdiv(N, meta['BLOCK_N'])) + _qkv_proj_split_kernel[grid]( + x, w_qkv, bias_arg, + Q, Kt, V, + M, K, N, q_size, kv_size, + x.stride(0), x.stride(1), + w_qkv.stride(0), w_qkv.stride(1), + Q.stride(0), Q.stride(1), + Kt.stride(0), Kt.stride(1), + V.stride(0), V.stride(1), + HAS_BIAS=HAS_BIAS, + ) + return Q, Kt, V + + +# ────────────────────────────────────────────────────────────────────────────── +# Pass 2 — RoPE for Q and K (standard rotate_half / contiguous-half convention) +# ────────────────────────────────────────────────────────────────────────────── +# cos/sin are [M, head_dim] (already gathered per-row by position_id upstream). +# rotate_half: for head split into [x1 | x2] (each head_dim/2 wide), +# out = [x1*cos1 - x2*sin1 , x2*cos2 + x1*sin2] +# with cos = [cos1 | cos2], sin = [sin1 | sin2] and cos1==cos2, sin1==sin2 in the +# standard HF layout. This matches torch: +# (t * cos) + (rotate_half(t) * sin), rotate_half(t) = cat(-t2, t1) + +@triton.jit +def _rope_inplace_kernel( + t_ptr, # [M, n_heads * head_dim] + cos_ptr, sin_ptr, # [M, head_dim] + M, n_heads, + head_dim: tl.constexpr, + half: tl.constexpr, + stride_tm, stride_td, + stride_cm, stride_cd, + BLOCK_H: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_h = tl.program_id(1) # head index + if pid_m >= M: + return + + h = pid_h + if h >= n_heads: + return + + half_idx = tl.arange(0, half) + base = pid_m * stride_tm + h * head_dim * stride_td + + # x1 = first half, x2 = second half + x1 = tl.load(t_ptr + base + half_idx * stride_td) + x2 = tl.load(t_ptr + base + (half + half_idx) * stride_td) + + cos1 = tl.load(cos_ptr + pid_m * stride_cm + half_idx * stride_cd) + sin1 = tl.load(sin_ptr + pid_m * stride_cm + half_idx * stride_cd) + cos2 = tl.load(cos_ptr + pid_m * stride_cm + (half + half_idx) * stride_cd) + sin2 = tl.load(sin_ptr + pid_m * stride_cm + (half + half_idx) * stride_cd) + + x1f = x1.to(tl.float32) + x2f = x2.to(tl.float32) + out1 = x1f * cos1.to(tl.float32) - x2f * sin1.to(tl.float32) + out2 = x2f * cos2.to(tl.float32) + x1f * sin2.to(tl.float32) + + tl.store(t_ptr + base + half_idx * stride_td, out1.to(t_ptr.dtype.element_ty)) + tl.store(t_ptr + base + (half + half_idx) * stride_td, out2.to(t_ptr.dtype.element_ty)) + + +def apply_rope_qk( + t: torch.Tensor, # [M, n_heads * head_dim] BF16, modified in place + cos: torch.Tensor, # [M, head_dim] + sin: torch.Tensor, # [M, head_dim] + head_dim: int, +): + """Apply standard rotate_half RoPE in place to a [M, n_heads*head_dim] tensor.""" + M, total = t.shape + assert total % head_dim == 0 + n_heads = total // head_dim + half = head_dim // 2 + grid = (M, n_heads) + _rope_inplace_kernel[grid]( + t, cos, sin, + M, n_heads, head_dim, half, + t.stride(0), t.stride(1), + cos.stride(0), cos.stride(1), + BLOCK_H=1, + ) + return t + + +def qkv_proj_rope( + x: torch.Tensor, # [M, K] BF16 + w_qkv: torch.Tensor, # [N, K] BF16 + bias: torch.Tensor | None, + q_size: int, + kv_size: int, + head_dim: int, + rope_cos: torch.Tensor | None, # [M, head_dim] + rope_sin: torch.Tensor | None, # [M, head_dim] +): + """QKV projection + split, with optional standard rotate_half RoPE on Q and K. + + Returns (Q, K, V), each BF16. RoPE is applied to Q and K only when + rope_cos/rope_sin are provided. V is never rotated. + """ + Q, Kt, V = qkv_proj_split(x, w_qkv, bias, q_size, kv_size) + if rope_cos is not None and rope_sin is not None: + apply_rope_qk(Q, rope_cos, rope_sin, head_dim) + apply_rope_qk(Kt, rope_cos, rope_sin, head_dim) + return Q, Kt, V diff --git a/docs/BLACKWELL_KERNELS_WIP.md b/docs/BLACKWELL_KERNELS_WIP.md new file mode 100644 index 000000000..30f13a264 --- /dev/null +++ b/docs/BLACKWELL_KERNELS_WIP.md @@ -0,0 +1 @@ +# WIP: [wip] blackwell: sm100 QKV+RoPE Triton kernel stub