diff --git a/aiter/ops/triton/_triton_kernels/pa_mqa_logits.py b/aiter/ops/triton/_triton_kernels/pa_mqa_logits.py index 86c2aa88a7..3c474c5a79 100644 --- a/aiter/ops/triton/_triton_kernels/pa_mqa_logits.py +++ b/aiter/ops/triton/_triton_kernels/pa_mqa_logits.py @@ -198,6 +198,7 @@ def _deepgemm_fp8_paged_mqa_logits_ragged_k( + (pid_batch * next_n + pid_next_n) * stride_out_batch + (context_idx + tl.arange(0, ChunkK)), logits, + mask=(context_idx + tl.arange(0, ChunkK)) < max_model_len, ) @@ -283,7 +284,9 @@ def _deepgemm_fp8_paged_mqa_logits_stage1( o = tl.maximum(o, 0.0) o = o * scale_weight[None, :].T - mask = context_idx + tl.arange(0, ChunkK) <= context_length - pid_next_n + mask = ( + context_idx + tl.arange(0, ChunkK) <= context_length - next_n + pid_next_n + ) o = tl.where(mask[None, :], o, float("-inf")) tl.store( @@ -377,7 +380,9 @@ def _deepgemm_fp8_paged_mqa_logits( o = tl.maximum(o, 0.0) o = o * scale_weight[None, :].T - mask = context_idx + tl.arange(0, ChunkK) <= context_length - pid_next_n + mask = ( + context_idx + tl.arange(0, ChunkK) <= context_length - next_n + pid_next_n + ) o = tl.where(mask[None, :], o, float("-inf")) logits = tl.reduce(o, axis=0, combine_fn=_sum_combine) @@ -386,4 +391,69 @@ def _deepgemm_fp8_paged_mqa_logits( + (pid_batch * next_n + pid_next_n) * stride_out_batch + (context_idx + tl.arange(0, ChunkK)), logits, + mask=(context_idx + tl.arange(0, ChunkK)) < max_model_len, ) + + +@triton.jit +def _gluon_deepgemm_fp8_paged_mqa_logits( + batch_size, + next_n, + heads_num, + Q_buffer, + stride_q_batch, + stride_q_next_n, + stride_q_heads, + KV_buffer, + stride_k_seq, + scale_buffer, + stride_scale_seq, + context_len_ptr, + kv_indices, + weights, + stride_w_batch, + OutLogits_buffer, + stride_out_batch, + max_model_len, + max_block_len, + SplitKV, + dummyPointerArg, + ChunkQ: tl.constexpr, + ChunkK: tl.constexpr, + HiddenDim: tl.constexpr, + KVBlockSize: tl.constexpr = 1, +): + # for AOT load use, only need kernel have the same signature as implementation side + pass + + +@triton.jit +def _gluon_deepgemm_fp8_paged_mqa_logits_preshuffle( + batch_size, + next_n, + heads_num, + Q_buffer, + stride_q_batch, + stride_q_next_n, + stride_q_heads, + KV_buffer, + stride_k_seq, + scale_buffer, + stride_scale_seq, + context_len_ptr, + kv_indices, + weights, + stride_w_batch, + OutLogits_buffer, + stride_out_batch, + max_model_len, + max_block_len, + SplitKV, + dummyPointerArg, + ChunkQ: tl.constexpr, + ChunkK: tl.constexpr, + HiddenDim: tl.constexpr, + KVBlockSize: tl.constexpr = 16, +): + # for AOT load use, only need kernel have the same signature as implementation side + pass diff --git a/aiter/ops/triton/gluon/pa_mqa_logits.py b/aiter/ops/triton/gluon/pa_mqa_logits.py new file mode 100644 index 0000000000..16129203d5 --- /dev/null +++ b/aiter/ops/triton/gluon/pa_mqa_logits.py @@ -0,0 +1,710 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +import math +import triton +import triton.language as tl + +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +try: + from triton.experimental.gluon.language.amd.cdna3 import ( + sched_barrier as _amd_iglp_sched_barrier, + ) + from triton.experimental.gluon.language.amd.cdna3 import ( + sched_group_barrier as _amd_iglp_sched_group_barrier, + ) +except ImportError: + # ignore iglp hint + @gluon.jit + def _amd_iglp_sched_barrier(inst_mask): + pass + + @gluon.jit + def _amd_iglp_sched_group_barrier(inst_mask, cnt, _): + pass + + +@triton.jit +def _sum_combine(a, b): + return a + b + + +@gluon.jit +def _gluon_deepgemm_fp8_paged_mqa_logits( + batch_size, + next_n, + heads_num, + Q_buffer, + stride_q_batch, + stride_q_next_n, + stride_q_heads, + KV_buffer, + stride_k_seq, + scale_buffer, + stride_scale_seq, + context_len_ptr, + kv_indices, + weights, + stride_w_batch, + OutLogits_buffer, + stride_out_batch, + max_model_len, + max_block_len, + SplitKV, + ChunkQ: tl.constexpr, + ChunkK: tl.constexpr, + HiddenDim: tl.constexpr, + KVBlockSize: tl.constexpr = 1, +): + pid = tl.program_id(0) + num_block_q_head = tl.cdiv(heads_num, ChunkQ) + + pid_q_head, remain_pid = pid % num_block_q_head, pid // num_block_q_head + pid_next_n, remain_pid = remain_pid % next_n, remain_pid // next_n + pid_batch, pid_split_kv = remain_pid % batch_size, remain_pid // batch_size + + context_length = gl.load(context_len_ptr + pid_batch) + + context_chunk_num = tl.cdiv(context_length, ChunkK) + split_context_chunk_num = tl.cdiv(context_chunk_num, SplitKV) + + split_context_start = (pid_split_kv * split_context_chunk_num) * ChunkK + split_context_length = min( + context_length - split_context_start, split_context_chunk_num * ChunkK + ) + + if split_context_length <= 0: + return + + residual_context = (ChunkK - split_context_length % ChunkK) % ChunkK + + NumWarps: gl.constexpr = 4 + ThreadsPerWarp: gl.constexpr = 64 + + # ===--------------------------------------------------- + # Gluon Layout + # ===--------------------------------------------------- + ValQMPerThread: gl.constexpr = ChunkQ // ( + NumWarps * ThreadsPerWarp // (HiddenDim // 16) + ) + layout_q: gl.constexpr = gl.BlockedLayout( + size_per_thread=[ValQMPerThread, 16], # q type is fp8 (E4M3) + threads_per_warp=[ThreadsPerWarp // (HiddenDim // 16), HiddenDim // 16], + warps_per_cta=[NumWarps, 1], + order=[1, 0], + ) + + ValKNPerThread: gl.constexpr = ChunkK // ( + NumWarps * ThreadsPerWarp // (HiddenDim // 16) + ) + layout_kv: gl.constexpr = gl.BlockedLayout( + size_per_thread=[ValKNPerThread, 16], # k type is fp8 (E4M3) + threads_per_warp=[ThreadsPerWarp // (HiddenDim // 16), HiddenDim // 16], + warps_per_cta=[NumWarps, 1], + order=[1, 0], + ) + + mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=3, + instr_shape=[16, 16], + transposed=False, + warps_per_cta=[1, NumWarps], + ) + mfma_layout_a: gl.constexpr = gl.DotOperandLayout( + operand_index=0, parent=mfma_layout, k_width=16 + ) + mfma_layout_b: gl.constexpr = gl.DotOperandLayout( + operand_index=1, parent=mfma_layout, k_width=16 + ) + + layout_scale: gl.constexpr = gl.SliceLayout(1, mfma_layout) + + # ===--------------------------------------------------- + # Pipeline Start + # ===--------------------------------------------------- + q = gl.amd.cdna3.buffer_load( + ptr=Q_buffer, + offsets=pid_batch * stride_q_batch + + pid_next_n * stride_q_next_n + + ( + ( + pid_q_head * ChunkQ + + gl.arange(0, ChunkQ, layout=gl.SliceLayout(1, layout_q)) + ) + * stride_q_heads + )[:, None] + + gl.arange(0, HiddenDim, layout=gl.SliceLayout(0, layout_q))[None, :], + ) + scale_weight = gl.amd.cdna3.buffer_load( + ptr=weights, + offsets=(pid_batch * next_n + pid_next_n) * stride_w_batch + + pid_q_head * ChunkQ + + gl.arange(0, ChunkQ, layout=layout_scale), + ) + + mask_kv_next = ( + split_context_start + - residual_context + + gl.arange(0, ChunkK, layout=gl.SliceLayout(1, layout_kv)) + >= 0 + ) + mask_kv_scale_next = ( + split_context_start + - residual_context + + gl.arange(0, ChunkK, layout=gl.SliceLayout(0, mfma_layout)) + >= 0 + ) + context_kv_idx_next = gl.amd.cdna3.buffer_load( + ptr=kv_indices, + offsets=pid_batch * max_block_len + + split_context_start + - residual_context + + gl.arange(0, ChunkK, layout=gl.SliceLayout(1, layout_kv)), + mask=mask_kv_next, + ) + context_kv_scale_idx_next = gl.amd.cdna3.buffer_load( + ptr=kv_indices, + offsets=pid_batch * max_block_len + + split_context_start + - residual_context + + gl.arange(0, ChunkK, layout=gl.SliceLayout(0, mfma_layout)), + mask=mask_kv_scale_next, + ) + + mfma_q = gl.convert_layout(q, mfma_layout_a) + + context_kv_idx_next = tl.where(mask_kv_next, context_kv_idx_next, 0) + k_next = gl.amd.cdna3.buffer_load( + ptr=KV_buffer, + offsets=context_kv_idx_next[:, None] * stride_k_seq + + gl.arange(0, HiddenDim, layout=gl.SliceLayout(0, layout_kv))[None, :], + ) + context_kv_scale_idx_next = tl.where( + mask_kv_scale_next, context_kv_scale_idx_next, 0 + ) + k_scale_f_next = gl.amd.cdna3.buffer_load( + ptr=scale_buffer, offsets=context_kv_scale_idx_next * stride_scale_seq + ) + + zero = gl.zeros((ChunkQ, ChunkK), dtype=tl.float32, layout=mfma_layout) + for context_idx in range( + split_context_start - residual_context, + split_context_start + split_context_length - ChunkK, + ChunkK, + ): + k = k_next + k_scale_f = k_scale_f_next + + context_kv_idx_next = gl.amd.cdna3.buffer_load( + ptr=kv_indices, + offsets=pid_batch * max_block_len + + context_idx + + ChunkK + + gl.arange(0, ChunkK, layout=gl.SliceLayout(1, layout_kv)), + ) + context_kv_scale_idx_next = gl.amd.cdna3.buffer_load( + ptr=kv_indices, + offsets=pid_batch * max_block_len + + context_idx + + ChunkK + + gl.arange(0, ChunkK, layout=gl.SliceLayout(0, mfma_layout)), + ) + + #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + #!=---------------------------- + mfma_k = gl.convert_layout(k.T, mfma_layout_b) + + o = gl.amd.cdna3.mfma(mfma_q, mfma_k, zero) + o = o * k_scale_f[None, :] + + #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + #!=---------------------------- + k_next = gl.amd.cdna3.buffer_load( + ptr=KV_buffer, + offsets=context_kv_idx_next[:, None] * stride_k_seq + + gl.arange(0, HiddenDim, layout=gl.SliceLayout(0, layout_kv))[None, :], + ) + o = gl.maximum(o, 0.0) + o = o * scale_weight[:, None] + + #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + #!=---------------------------- + k_scale_f_next = gl.amd.cdna3.buffer_load( + ptr=scale_buffer, offsets=context_kv_scale_idx_next * stride_scale_seq + ) + + mask = ( + context_idx + gl.arange(0, ChunkK, layout=gl.SliceLayout(0, mfma_layout)) + <= context_length - next_n + pid_next_n + ) + o = tl.where(mask[None, :], o, float("-inf")) + + logits = gl.reduce(o, axis=0, combine_fn=_sum_combine) + gl.amd.cdna3.buffer_store( + logits, + ptr=OutLogits_buffer, + offsets=(pid_batch * next_n + pid_next_n) * stride_out_batch + + ( + context_idx + + gl.arange(0, ChunkK, layout=gl.SliceLayout(0, mfma_layout)) + ), + mask=context_idx + + gl.arange(0, ChunkK, layout=gl.SliceLayout(0, mfma_layout)) + >= 0, + ) + + context_idx = split_context_start + split_context_length - ChunkK + k = k_next + k_scale_f = k_scale_f_next + + mfma_k = gl.convert_layout(k.T, mfma_layout_b) + o = gl.amd.cdna3.mfma(mfma_q, mfma_k, zero) + + o = o * k_scale_f[None, :] + o = gl.maximum(o, 0.0) + o = o * scale_weight[:, None] + + mask = ( + context_idx + gl.arange(0, ChunkK, layout=gl.SliceLayout(0, mfma_layout)) + <= context_length - next_n + pid_next_n + ) + o = tl.where(mask[None, :], o, float("-inf")) + + logits = gl.reduce(o, axis=0, combine_fn=_sum_combine) + gl.amd.cdna3.buffer_store( + logits, + ptr=OutLogits_buffer, + offsets=(pid_batch * next_n + pid_next_n) * stride_out_batch + + (context_idx + gl.arange(0, ChunkK, layout=gl.SliceLayout(0, mfma_layout))), + mask=context_idx + gl.arange(0, ChunkK, layout=gl.SliceLayout(0, mfma_layout)) + >= 0, + ) + + +@gluon.jit +def _gluon_deepgemm_fp8_paged_mqa_logits_preshuffle( + batch_size, + next_n, + heads_num, + Q_buffer, + stride_q_batch, + stride_q_next_n, + stride_q_heads, + KV_buffer, + stride_k_seq, + scale_buffer, + stride_scale_seq, + context_len_ptr, + kv_indices, + weights, + stride_w_batch, + OutLogits_buffer, + stride_out_batch, + max_model_len, + max_block_len, + SplitKV, + ChunkQ: tl.constexpr, + ChunkK: tl.constexpr, + HiddenDim: tl.constexpr, + KVBlockSize: tl.constexpr = 16, +): + # ===--------------------------------------------------- + # Gluon Layout + # ===--------------------------------------------------- + NumWarps: gl.constexpr = 4 + ThreadsPerWarp: gl.constexpr = 64 + + ValQMPerThread: gl.constexpr = ChunkQ // ( + NumWarps * ThreadsPerWarp // (HiddenDim // 16) + ) + layout_q: gl.constexpr = gl.BlockedLayout( + size_per_thread=[ValQMPerThread, 16], # q type is fp8 (E4M3) + threads_per_warp=[ThreadsPerWarp // (HiddenDim // 16), HiddenDim // 16], + warps_per_cta=[NumWarps, 1], + order=[1, 0], + ) + + ChunkKPerStage: gl.constexpr = ChunkK // 2 + MFMAPerWarp: gl.constexpr = ChunkKPerStage // 16 // NumWarps + + mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=3, + instr_shape=[16, 16], + transposed=False, + warps_per_cta=[1, NumWarps], + tiles_per_warp=[1, MFMAPerWarp], + ) + mfma_layout_a: gl.constexpr = gl.DotOperandLayout( + operand_index=0, parent=mfma_layout, k_width=16 + ) + mfma_layout_b: gl.constexpr = gl.DotOperandLayout( + operand_index=1, parent=mfma_layout, k_width=16 + ) + + layout_scale: gl.constexpr = gl.SliceLayout(1, mfma_layout) + + ContextBlockPerChunkK: gl.constexpr = ChunkK // KVBlockSize + + DS_WRITE: gl.constexpr = 0x200 + DS_READ: gl.constexpr = 0x100 + BUFFER_LOAD: gl.constexpr = 0x020 + MFMA: gl.constexpr = 0x008 + VALU: gl.constexpr = 0x002 + + # ===--------------------------------------------------- + # Mapping WorkTile + # ===--------------------------------------------------- + pid = tl.program_id(0) + + # ===--------------------------------------------------- + pid_batch, remain_pid = pid % batch_size, pid // batch_size + pid_next_n, pid_split_kv = remain_pid % next_n, remain_pid // next_n + # ===--------------------------------------------------- + context_length = gl.load(context_len_ptr + pid_batch) + + context_chunk_num = tl.cdiv(context_length, ChunkK) + split_context_chunk_num = context_chunk_num // SplitKV + residual_context_chunks = context_chunk_num % SplitKV + split_context_start = ( + pid_split_kv * split_context_chunk_num * ChunkK + + min(pid_split_kv, residual_context_chunks) * ChunkK + ) + split_context_length = min( + context_length - split_context_start, + split_context_chunk_num * ChunkK + + (ChunkK if pid_split_kv < residual_context_chunks else 0), + ) + + if split_context_length <= 0: + return + + split_context_block = tl.cdiv(split_context_length, KVBlockSize) + split_context_length = split_context_block * KVBlockSize + + residual_context_blocks = ( + ContextBlockPerChunkK - split_context_block % ContextBlockPerChunkK + ) % ContextBlockPerChunkK + residual_context = residual_context_blocks * KVBlockSize + + # ===--------------------------------------------------- + # Pipeline Start + _amd_iglp_sched_barrier(0x0) + # ===--------------------------------------------------- + q = gl.amd.cdna3.buffer_load( + ptr=Q_buffer, + offsets=pid_batch * stride_q_batch + + pid_next_n * stride_q_next_n + + (gl.arange(0, ChunkQ, layout=gl.SliceLayout(1, layout_q)) * stride_q_heads)[ + :, None + ] + + gl.arange(0, HiddenDim, layout=gl.SliceLayout(0, layout_q))[None, :], + ) + + context_idx = split_context_start - residual_context + + mask_kv_next_0 = ( + context_idx // KVBlockSize + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + // KVBlockSize + ) >= split_context_start // KVBlockSize + context_kv_idx_next_0 = gl.amd.cdna3.buffer_load( + ptr=kv_indices, + offsets=pid_batch * max_block_len + + context_idx // KVBlockSize + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + // KVBlockSize, + mask=mask_kv_next_0, + ) + + mask_kv_next_1 = ( + (context_idx + ChunkKPerStage) // KVBlockSize + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + // KVBlockSize + ) >= split_context_start // KVBlockSize + context_kv_idx_next_1 = gl.amd.cdna3.buffer_load( + ptr=kv_indices, + offsets=pid_batch * max_block_len + + (context_idx + ChunkKPerStage) // KVBlockSize + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + // KVBlockSize, + mask=mask_kv_next_1, + ) + + scale_weight = gl.amd.cdna3.buffer_load( + ptr=weights, + offsets=(pid_batch * next_n + pid_next_n) * stride_w_batch + + gl.arange(0, ChunkQ, layout=layout_scale), + ) + + offset_k_fixed = ( + gl.arange(0, HiddenDim, layout=gl.SliceLayout(1, mfma_layout_b)) % 16 + + gl.arange(0, HiddenDim, layout=gl.SliceLayout(1, mfma_layout_b)) // 16 * 256 + )[:, None] + ( + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) % 16 * 16 + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + % KVBlockSize + // 16 + * 16 + * 128 + )[ + None, : + ] + + #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + #!=---------------------------- + mfma_q = gl.convert_layout(q, mfma_layout_a) + + context_kv_idx_next_0 = tl.where(mask_kv_next_0, context_kv_idx_next_0, 0) + k_next_0 = gl.amd.cdna3.buffer_load( + ptr=KV_buffer, + offsets=offset_k_fixed + context_kv_idx_next_0[None, :] * stride_k_seq, + ) + k_scale_f_next_0 = gl.amd.cdna3.buffer_load( + ptr=scale_buffer, + offsets=context_kv_idx_next_0 * stride_scale_seq + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + % KVBlockSize, + ) + + _amd_iglp_sched_group_barrier(DS_READ, 4, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 4, 0) + _amd_iglp_sched_group_barrier(DS_READ, 2, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(DS_READ, 2, 0) + + if context_idx + ChunkK < split_context_start + split_context_length: + context_kv_idx_next_0 = gl.amd.cdna3.buffer_load( + ptr=kv_indices, + offsets=pid_batch * max_block_len + + (context_idx + ChunkK) // KVBlockSize + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + // KVBlockSize, + ) + #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + #!=---------------------------- + + # ===--------------------------------------------------- + # Precompute First Iteration + # ===--------------------------------------------------- + zero = gl.zeros((ChunkQ, ChunkKPerStage), dtype=tl.float32, layout=mfma_layout) + + k = k_next_0 + k_scale_f = k_scale_f_next_0 + + #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + #!=---------------------------- + + context_kv_idx_next_1 = tl.where(mask_kv_next_1, context_kv_idx_next_1, 0) + k_next_1 = gl.amd.cdna3.buffer_load( + ptr=KV_buffer, + offsets=offset_k_fixed + context_kv_idx_next_1[None, :] * stride_k_seq, + ) + k_scale_f_next_1 = gl.amd.cdna3.buffer_load( + ptr=scale_buffer, + offsets=context_kv_idx_next_1 * stride_scale_seq + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + % KVBlockSize, + ) + mfma_k = gl.convert_layout(k, mfma_layout_b) + o = gl.amd.cdna3.mfma(mfma_q, mfma_k, zero) + + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + #!=---------------------------- + + k_scale_f = gl.convert_layout(k_scale_f, gl.SliceLayout(0, mfma_layout)) + + o = o * k_scale_f[None, :] + o = gl.maximum(o, 0.0) + o = o * scale_weight[:, None] + + logits = gl.reduce(o, axis=0, combine_fn=_sum_combine) + gl.amd.cdna3.buffer_store( + logits, + ptr=OutLogits_buffer, + offsets=(pid_batch * next_n + pid_next_n) * stride_out_batch + + ( + context_idx + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout)) + ), + mask=context_idx + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout)) + >= split_context_start, + ) + + for context_idx in range( + split_context_start - residual_context, + split_context_start + split_context_length - ChunkK, + ChunkK, + ): + k = k_next_1 + k_scale_f = k_scale_f_next_1 + + #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + #!=---------------------------- + + context_kv_idx_next_1 = gl.amd.cdna3.buffer_load( + ptr=kv_indices, + offsets=pid_batch * max_block_len + + (context_idx + ChunkK + ChunkKPerStage) // KVBlockSize + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + // KVBlockSize, + ) + k_next_0 = gl.amd.cdna3.buffer_load( + ptr=KV_buffer, + offsets=offset_k_fixed + context_kv_idx_next_0[None, :] * stride_k_seq, + ) + k_scale_f_next_0 = gl.amd.cdna3.buffer_load( + ptr=scale_buffer, + offsets=context_kv_idx_next_0 * stride_scale_seq + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + % KVBlockSize, + ) + mfma_k = gl.convert_layout(k, mfma_layout_b) + o = gl.amd.cdna3.mfma(mfma_q, mfma_k, zero) + + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + #!=---------------------------- + k_scale_f = gl.convert_layout(k_scale_f, gl.SliceLayout(0, mfma_layout)) + o = o * k_scale_f[None, :] + o = gl.maximum(o, 0.0) + o = o * scale_weight[:, None] + + logits = gl.reduce(o, axis=0, combine_fn=_sum_combine) + gl.amd.cdna3.buffer_store( + logits, + ptr=OutLogits_buffer, + offsets=(pid_batch * next_n + pid_next_n) * stride_out_batch + + ( + context_idx + + ChunkKPerStage + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout)) + ), + mask=context_idx + + ChunkKPerStage + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout)) + >= split_context_start, + ) + + # ======================================================================================= + + k = k_next_0 + k_scale_f = k_scale_f_next_0 + + # #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + # #!=---------------------------- + if context_idx + ChunkK + ChunkK < split_context_start + split_context_length: + context_kv_idx_next_0 = gl.amd.cdna3.buffer_load( + ptr=kv_indices, + offsets=pid_batch * max_block_len + + (context_idx + ChunkK + ChunkK) // KVBlockSize + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + // KVBlockSize, + ) + k_next_1 = gl.amd.cdna3.buffer_load( + ptr=KV_buffer, + offsets=offset_k_fixed + context_kv_idx_next_1[None, :] * stride_k_seq, + ) + k_scale_f_next_1 = gl.amd.cdna3.buffer_load( + ptr=scale_buffer, + offsets=context_kv_idx_next_1 * stride_scale_seq + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout_b)) + % KVBlockSize, + ) + mfma_k = gl.convert_layout(k, mfma_layout_b) + o = gl.amd.cdna3.mfma(mfma_q, mfma_k, zero) + + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + _amd_iglp_sched_group_barrier(BUFFER_LOAD, 2, 0) + _amd_iglp_sched_group_barrier(MFMA, 8, 0) + #!=---------------------------- + _amd_iglp_sched_barrier(0x0) + #!=---------------------------- + + k_scale_f = gl.convert_layout(k_scale_f, gl.SliceLayout(0, mfma_layout)) + + o = o * k_scale_f[None, :] + o = gl.maximum(o, 0.0) + o = o * scale_weight[:, None] + + logits = gl.reduce(o, axis=0, combine_fn=_sum_combine) + + gl.amd.cdna3.buffer_store( + logits, + ptr=OutLogits_buffer, + offsets=(pid_batch * next_n + pid_next_n) * stride_out_batch + + ( + context_idx + + ChunkK + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout)) + ), + ) + + context_idx = split_context_start + split_context_length - ChunkK + + k = k_next_1 + k_scale_f = k_scale_f_next_1 + + mfma_k = gl.convert_layout(k, mfma_layout_b) + o = gl.amd.cdna3.mfma(mfma_q, mfma_k, zero) + k_scale_f = gl.convert_layout(k_scale_f, gl.SliceLayout(0, mfma_layout)) + o = o * k_scale_f[None, :] + o = gl.maximum(o, 0.0) + o = o * scale_weight[:, None] + + mask = ( + context_idx + + ChunkKPerStage + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout)) + <= context_length - next_n + pid_next_n + ) + o = tl.where(mask[None, :], o, float("-inf")) + + logits = gl.reduce(o, axis=0, combine_fn=_sum_combine) + gl.amd.cdna3.buffer_store( + logits, + ptr=OutLogits_buffer, + offsets=(pid_batch * next_n + pid_next_n) * stride_out_batch + + ( + context_idx + + ChunkKPerStage + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout)) + ), + mask=context_idx + + ChunkKPerStage + + gl.arange(0, ChunkKPerStage, layout=gl.SliceLayout(0, mfma_layout)) + >= split_context_start, + ) diff --git a/aiter/ops/triton/pa_mqa_logits.py b/aiter/ops/triton/pa_mqa_logits.py index d343eeed6e..0c6ecceb8b 100644 --- a/aiter/ops/triton/pa_mqa_logits.py +++ b/aiter/ops/triton/pa_mqa_logits.py @@ -1,15 +1,74 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. +# ======================================================================== +# How to use AOT gluon kernel for pa_mqa_logits on lower triton version (below 3.4.0): +# 1. Generate Gluon kernel based on rocm/triton/gluon_ext (3.5.0+gite392a058) +# it requires zip installed. +# $ cd ${AOT_DUMP_AITER_ROOT} +# $ python3 op_tests/op_benchmarks/triton/bench_deepgemm_attention.py --batch=1 -aot [-p] +# "-p" means kernel could assume the stride of KVCache is aligned to 16B. +# If enable it, the stride of KVCache in the AOT_load side must also be aligned to 16B. +# 2. Copy generated paged_mqa_logits_aot_kernel.zip to ${AOT_LOAD_AITER_ROOT}/aiter/ops/triton/configs +# and unzip it. +# $ cd ${AOT_LOAD_AITER_ROOT} +# $ cd aiter/ops/triton/configs && unzip paged_mqa_logits_aot_kernel.zip && cd - +# 3. Set env variable to enable AOT gluon kernel loading +# $ export AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=1 +# $ python3 op_tests/op_benchmarks/triton/bench_deepgemm_attention.py -kv_length=32768 --batch=2 -mtp=1 -p +# Set AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS=0 to disable AOT gluon kernel. It will backward +# to triton JIT kernel +# ======================================================================== + +import os import torch +import triton +from functools import lru_cache + +from triton.backends.compiler import GPUTarget -from aiter.ops.triton._triton_kernels.pa_mqa_logits import ( - _deepgemm_fp8_paged_mqa_logits_stage1, - _deepgemm_fp8_paged_mqa_logits_stage1_ragged_k, - _deepgemm_fp8_paged_mqa_logits, - _deepgemm_fp8_paged_mqa_logits_ragged_k, +enable_aot_gluon_pa_mqa_logits = os.environ.get( + "AITER_ENABLE_AOT_GLUON_PA_MQA_LOGITS", "0" +) +enable_aot_gluon_pa_mqa_logits = enable_aot_gluon_pa_mqa_logits == "1" + +if triton.__version__ >= "3.5.0": + from triton.experimental.gluon._runtime import GluonASTSource as ASTSource + from aiter.ops.triton._triton_kernels.pa_mqa_logits import ( + _deepgemm_fp8_paged_mqa_logits_stage1, + _deepgemm_fp8_paged_mqa_logits_stage1_ragged_k, + _deepgemm_fp8_paged_mqa_logits, + _deepgemm_fp8_paged_mqa_logits_ragged_k, + ) + from aiter.ops.triton.gluon.pa_mqa_logits import ( + _gluon_deepgemm_fp8_paged_mqa_logits, + _gluon_deepgemm_fp8_paged_mqa_logits_preshuffle, + ) + + enable_gluon_pa_mqa_logits = True + enable_jit_gluon_pa_mqa_logits_kernel = True +else: + from triton.compiler import ASTSource + from aiter.ops.triton._triton_kernels.pa_mqa_logits import ( + _deepgemm_fp8_paged_mqa_logits_stage1, + _deepgemm_fp8_paged_mqa_logits_stage1_ragged_k, + _deepgemm_fp8_paged_mqa_logits, + _deepgemm_fp8_paged_mqa_logits_ragged_k, + _gluon_deepgemm_fp8_paged_mqa_logits, + _gluon_deepgemm_fp8_paged_mqa_logits_preshuffle, + ) + + assert triton.__version__ < "3.4.0" + enable_gluon_pa_mqa_logits = enable_aot_gluon_pa_mqa_logits + enable_jit_gluon_pa_mqa_logits_kernel = False + + +from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH +from aiter.utility.triton.triton_metadata_redirect import ( + AOTMetadataContext, ) from aiter import dtypes +from ...jit.utils.chip_info import get_gfx def deepgemm_fp8_paged_mqa_logits_ragged_k( @@ -122,9 +181,17 @@ def deepgemm_fp8_paged_mqa_logits_stage1( context_lens: torch.Tensor, kv_indices: torch.Tensor, max_model_len: int, + ChunkQ: int = 64, + ChunkK: int = 256, + TotalCuCount: int = 80, + WavePerEU: int = 2, ): batch_size, next_n, heads, hidden_dim = q_fp8.size() _, max_blk_len = kv_indices.size() + + TileQCount = batch_size * next_n * (heads // ChunkQ) + SplitKV = (max(1, TotalCuCount // TileQCount) + 4) // 5 * 5 * WavePerEU + kv_cache_fp8, kv_cache_scale = ( kv_cache_fp8[..., :hidden_dim], kv_cache_fp8[..., hidden_dim:], @@ -134,14 +201,14 @@ def deepgemm_fp8_paged_mqa_logits_stage1( kv_cache_fp8 = kv_cache_fp8.view(dtypes.fp8) config = { - "ChunkQ": 32, - "ChunkK": 64, + "ChunkQ": ChunkQ, + "ChunkK": ChunkK, "HiddenDim": hidden_dim, - "SplitKV": 5, + "SplitKV": SplitKV, } assert heads % config["ChunkQ"] == 0 - grid = (batch_size * next_n * (heads // config["ChunkQ"] * config["SplitKV"]),) + grid = (batch_size * next_n * (heads // config["ChunkQ"] * SplitKV),) _deepgemm_fp8_paged_mqa_logits_stage1[grid]( batch_size, next_n, @@ -163,58 +230,265 @@ def deepgemm_fp8_paged_mqa_logits_stage1( out_qk.stride(1), max_model_len, max_blk_len, + waves_per_eu=WavePerEU, **config, ) +@lru_cache(maxsize=None) +def _compile_deepgemm_fp8_paged_mqa_logits( + ChunkQ, + ChunkK, + Preshuffle, + KVBlockSize, + HiddenDim, + is_padded_mode: bool, + WavePerEU: int = 2, +): + gfx_version = get_gfx() + assert gfx_version == "gfx942" or gfx_version == "gfx950" + target = GPUTarget("hip", gfx_version, 64) + + gfx_fp8_pointer = "*fp8e4b8" if gfx_version == "gfx942" else "*fp8e4nv" + + fn_signature = { + "batch_size": "i32", + "next_n": "i32", + "heads_num": "i32", + "Q_buffer": gfx_fp8_pointer, + "stride_q_batch": "i32", + "stride_q_next_n": "i32", + "stride_q_heads": "i32", + "KV_buffer": gfx_fp8_pointer, + "stride_k_seq": "i32", + "scale_buffer": "*fp32", + "stride_scale_seq": "i32", + "context_len_ptr": "*i32", + "kv_indices": "*i32", + "weights": "*fp32", + "stride_w_batch": "i32", + "OutLogits_buffer": "*fp32", + "stride_out_batch": "i32", + "max_model_len": "i32", + "max_block_len": "i32", + "SplitKV": "i32", + } + if not enable_jit_gluon_pa_mqa_logits_kernel: + fn_signature["dummyPointerArg"] = "*i32" + fn_signature["ChunkQ"] = "constexpr" + fn_signature["ChunkK"] = "constexpr" + fn_signature["KVBlockSize"] = "constexpr" + fn_signature["HiddenDim"] = "constexpr" + + options = { + "num_warps": 4, + "waves_per_eu": WavePerEU, + "num_stages": 2, + "num_ctas": 1, + "cluster_dims": [1, 1, 1], + "arch": gfx_version, + "backend_name": "hip", + "warp_size": 64, + "name": ( + "_gluon_deepgemm_fp8_paged_mqa_logits" + if not Preshuffle + else "_gluon_deepgemm_fp8_paged_mqa_logits_preshuffle" + ), + } + + kv_cache_attr = [] + if is_padded_mode: + kv_cache_attr.append(["tt.divisibility", 16]) + + kernel_fn = ( + _gluon_deepgemm_fp8_paged_mqa_logits + if not Preshuffle + else _gluon_deepgemm_fp8_paged_mqa_logits_preshuffle + ) + src = ASTSource( + fn=kernel_fn, + signature=fn_signature, + constexprs={ + "ChunkQ": ChunkQ, + "ChunkK": ChunkK, + "KVBlockSize": KVBlockSize, + "HiddenDim": HiddenDim, + }, + attrs={ + (2,): [["tt.divisibility", 16]], # heads_num + (3,): [["tt.divisibility", 16], ["tt.pointer_range", 32]], # Q_buffer + (4,): [["tt.divisibility", 16]], # stride_q_batch + (5,): [["tt.divisibility", 16]], # stride_q_next_n + (6,): [["tt.divisibility", 16]], # stride_q_heads + (7,): kv_cache_attr, # KV_buffer + (8,): kv_cache_attr, # stride_k_seq + (9,): kv_cache_attr, # scale_buffer + (10,): kv_cache_attr, # stride_scale_seq + (11,): [["tt.pointer_range", 32]], # context_len_ptr + (12,): [["tt.pointer_range", 32]], # kv_indices + (13,): [ + ["tt.divisibility", 16], + ["tt.pointer_range", 32], + ], # weights + (14,): [["tt.divisibility", 16]], # stride_w_batch + (15,): [["tt.pointer_range", 32]], # OutLogits_buffer + }, + ) + + if enable_jit_gluon_pa_mqa_logits_kernel: + kernel = triton.compile( + src, + target=target, + options=options, + ) + else: + padded_str = "T" if is_padded_mode and not Preshuffle else "F" + kernel_str = f"paged_mqa_logits{"_preshuffle" if Preshuffle else ""}_{ChunkQ}x{ChunkK}x{HiddenDim}_B{KVBlockSize}P{padded_str}W{WavePerEU}" + metadata_pth = f"{AITER_TRITON_CONFIGS_PATH}/paged_mqa_logits/aot/{kernel_str}" + with AOTMetadataContext( + kernel_fn.fn.__name__, + metadata_pth, + ): + kernel = triton.compile( + src, + target=target, + options=options, + ) + return kernel + + def deepgemm_fp8_paged_mqa_logits( q_fp8: torch.Tensor, # dtype = float8 - kv_cache_fp8: torch.Tensor, # dtype = float8 [num_blocks, 1, 1, D+4] + kv_cache, weights: torch.Tensor, # dtype = float32 out_logits: torch.Tensor, # dtype = float32 context_lens: torch.Tensor, kv_indices: torch.Tensor, max_model_len: int, - ChunkK: int = 64, - SplitKV: int = 5, + Preshuffle: bool = False, + KVBlockSize: int = 1, + ChunkK: int = 256, + TotalCuCount: int = 80 if get_gfx() == "gfx942" else 256, + WavePerEU: int = 2, ): batch_size, next_n, heads, hidden_dim = q_fp8.size() - _, max_blk_len = kv_indices.size() + num_block, block_Size, _, index_dim = kv_cache.size() + _, max_block_len = kv_indices.size() + + TileQCount = batch_size * next_n + SplitKV = (max(1, TotalCuCount // TileQCount) + 4) // 5 * 5 * WavePerEU + + assert ChunkK % KVBlockSize == 0 + assert block_Size == KVBlockSize + if Preshuffle: + assert ( + KVBlockSize % 16 == 0 + ), f"Preshuffle mode only supports KVBlockSize aligned to 16. Got KVBlockSize={KVBlockSize}" + + kv_cache = kv_cache.view(-1, KVBlockSize * index_dim) kv_cache_fp8, kv_cache_scale = ( - kv_cache_fp8[..., :hidden_dim], - kv_cache_fp8[..., hidden_dim:], + kv_cache[..., : KVBlockSize * hidden_dim], + kv_cache[..., KVBlockSize * hidden_dim :], ) - # Since triton doesn't have the reinterpret_cast, we slice the scale out and view it as float - kv_cache_scale = kv_cache_scale.view(torch.float32) kv_cache_fp8 = kv_cache_fp8.view(dtypes.fp8) + kv_cache_scale = kv_cache_scale.view(torch.float32) - config = { - "ChunkQ": heads, - "ChunkK": ChunkK, - "HiddenDim": hidden_dim, - "SplitKV": SplitKV, - } - - grid = (batch_size * next_n * config["SplitKV"],) - _deepgemm_fp8_paged_mqa_logits[grid]( - batch_size, - next_n, - heads, - q_fp8, - q_fp8.stride(0), - q_fp8.stride(1), - q_fp8.stride(2), - kv_cache_fp8, - kv_cache_fp8.stride(0), - kv_cache_scale, - kv_cache_scale.stride(0), - context_lens, - kv_indices, - weights, - weights.stride(0), - out_logits, - out_logits.stride(0), - max_model_len, - max_blk_len, - **config, - ) + grid = (batch_size * next_n * SplitKV, 1, 1) + if enable_gluon_pa_mqa_logits: + is_padded_mode = kv_cache_fp8.stride(0) % 16 == 0 + kernel = _compile_deepgemm_fp8_paged_mqa_logits( + ChunkQ=heads, + ChunkK=ChunkK, + Preshuffle=Preshuffle, + KVBlockSize=KVBlockSize, + HiddenDim=hidden_dim, + is_padded_mode=is_padded_mode, + WavePerEU=WavePerEU, + ) + if enable_jit_gluon_pa_mqa_logits_kernel: + kernel[grid]( + batch_size, + next_n, + heads, + q_fp8, + q_fp8.stride(0), + q_fp8.stride(1), + q_fp8.stride(2), + kv_cache_fp8, + kv_cache_fp8.stride(0), + kv_cache_scale, + kv_cache_scale.stride(0), + context_lens, + kv_indices, + weights, + weights.stride(0), + out_logits, + out_logits.stride(0), + max_model_len, + max_block_len, + SplitKV, + # constexpr + heads, + ChunkK, + KVBlockSize, + hidden_dim, + ) + else: # load AOT compiled gluon kernel + kernel[grid]( + batch_size, + next_n, + heads, + q_fp8, + q_fp8.stride(0), + q_fp8.stride(1), + q_fp8.stride(2), + kv_cache_fp8, + kv_cache_fp8.stride(0), + kv_cache_scale, + kv_cache_scale.stride(0), + context_lens, + kv_indices, + weights, + weights.stride(0), + out_logits, + out_logits.stride(0), + max_model_len, + max_block_len, + SplitKV, + out_logits, # dummyPointerArg for triton version < 3.4.0, + # the kernel signature has an extra pointer argument on triton>=3.5.0 + # constexpr + heads, + ChunkK, + KVBlockSize, + hidden_dim, + ) + else: + assert not Preshuffle, "Preshuffle mode is only supported on gluon kernel." + kernel = _deepgemm_fp8_paged_mqa_logits[grid]( + batch_size, + next_n, + heads, + q_fp8, + q_fp8.stride(0), + q_fp8.stride(1), + q_fp8.stride(2), + kv_cache_fp8, + kv_cache_fp8.stride(0), + kv_cache_scale, + kv_cache_scale.stride(0), + context_lens, + kv_indices, + weights, + weights.stride(0), + out_logits, + out_logits.stride(0), + max_model_len, + max_block_len, + waves_per_eu=WavePerEU, + ChunkQ=heads, + ChunkK=ChunkK, + SplitKV=SplitKV, + HiddenDim=hidden_dim, + ) + return triton.runtime.cache.get_cache_manager(kernel.hash).key diff --git a/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py b/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py index 0b7a2589f4..5a6b66133a 100644 --- a/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py +++ b/op_tests/op_benchmarks/triton/bench_deepgemm_attention.py @@ -6,6 +6,7 @@ import pytest import torch +import os import triton import triton.language as tl @@ -17,28 +18,33 @@ deepgemm_fp8_paged_mqa_logits_stage1_ragged_k, deepgemm_fp8_paged_mqa_logits_ragged_k, ) +from aiter.ops.shuffle import shuffle_weight def cdiv(x: int, y: int) -> int: return (x + y - 1) // y -def kv_cache_cast_to_fp8(x: torch.Tensor) -> torch.Tensor: +def kv_cache_cast_to_fp8(x: torch.Tensor, padding=False) -> torch.Tensor: num_blocks, block_size, num_heads, head_dim = x.shape assert num_heads == 1 x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) sf = x_amax / 240.0 x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fnuz) + + padding_size = 0 if not padding else (16 - (block_size * 4) % 16) % 16 x_fp8 = torch.empty( - (num_blocks, block_size * (head_dim + 4)), device=x.device, dtype=torch.uint8 + (num_blocks, block_size * (head_dim + 4 + padding_size)), + device=x.device, + dtype=torch.uint8, ) x_fp8[:, : block_size * head_dim] = x_scaled.view( num_blocks, block_size * head_dim ).view(dtype=torch.uint8) - x_fp8[:, block_size * head_dim :] = sf.view(num_blocks, block_size).view( - dtype=torch.uint8 - ) - return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4) + x_fp8[:, block_size * head_dim : block_size * head_dim + 4 * block_size] = sf.view( + num_blocks, block_size + ).view(dtype=torch.uint8) + return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4 + padding_size) def ref_fp8_paged_mqa_logits( @@ -100,7 +106,7 @@ def ref_fp8_paged_mqa_logits_ragged( max_model_len: int, ): batch_size, next_n, heads, dim = q.size() - seq_kv, _, dim = kv_cache.size() # 3d + seq_kv, block_size, dim = kv_cache.size() # 3d logits = torch.full( [batch_size * next_n, max_model_len], float("-inf"), @@ -140,12 +146,29 @@ def ref_fp8_paged_mqa_logits_ragged( def create_paged_mqa_logits_configs(args: argparse.Namespace): x_names = ["batch_size", "next_n", "heads", "index_dim", "avg_kv_length"] - line_names = ["ragged_k", "non_ragged_k"] + line_names = ["non_ragged_k"] line_args = "kv_storage_kind" - x_vals_list = [ - (args.batch, args.mtp + 1, args.heads, args.index_dim, args.kv_length) - ] + if args.perf: + x_vals_list = [ + (1, 2, 64, 128, 16384), + (1, 2, 64, 128, 32768), + (1, 2, 64, 128, 65536), + (2, 2, 64, 128, 16384), + (2, 2, 64, 128, 32768), + (2, 2, 64, 128, 65536), + (4, 2, 64, 128, 16384), + (4, 2, 64, 128, 32768), + (4, 2, 64, 128, 65536), + (1, 1, 64, 128, 65536), + (2, 1, 64, 128, 65536), + (4, 1, 64, 128, 65536), + (8, 1, 64, 128, 65536), + ] + else: + x_vals_list = [ + (args.batch, args.mtp + 1, args.heads, args.index_dim, args.kv_length) + ] configs = [] configs.append( @@ -166,8 +189,8 @@ def create_paged_mqa_logits_configs(args: argparse.Namespace): def run_benchmark(args: argparse.Namespace): - ChunkK = 64 - SplitKV = 5 + ChunkK = 256 + WavePerEU = 2 @triton.testing.perf_report(create_paged_mqa_logits_configs(args)) def test_deepgemm_fp8_paged_mqa_logits( @@ -177,10 +200,12 @@ def test_deepgemm_fp8_paged_mqa_logits( random.seed(0) max_model_len = 2 * avg_kv_length - num_blocks = 111 * 1000 * 3 - blocksize = 1 + num_blocks = max_model_len + blocksize = args.blocksize if args.kv_preshuffle else 1 - var_ratio = 0.4 + assert blocksize == 1 or args.kv_preshuffle and blocksize % 16 == 0 + + var_ratio = 0.0 context_lens = ( torch.randint( int((1 - var_ratio) * avg_kv_length), @@ -229,7 +254,7 @@ def test_deepgemm_fp8_paged_mqa_logits( counter += 1 q_fp8 = q.to(qk_datatype) - kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache) + kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, padding=args.padding) kv_indices = torch.zeros( prefix_sum_context_lens[-1], device="cuda", dtype=torch.int32 @@ -247,19 +272,13 @@ def test_deepgemm_fp8_paged_mqa_logits( else: ref_logits = ref_fp8_paged_mqa_logits_ragged( q, - kv_cache.view([num_blocks, 1, index_dim]), + kv_cache.view([num_blocks, blocksize, index_dim]), weights, prefix_sum_context_lens, kv_indices, max_model_len, ) - out_qk = torch.full( - (heads, batch_size * next_n, max_model_len), - float("-inf"), - device="cuda", - dtype=torch.float32, - ) out_logits = torch.full( (batch_size * next_n, max_model_len), float("-inf"), @@ -268,15 +287,21 @@ def test_deepgemm_fp8_paged_mqa_logits( ) if kv_storage_kind == "non_ragged_k": - deepgemm_fp8_paged_mqa_logits_stage1( - q_fp8, - kv_cache_fp8, - weights, - out_qk, - context_lens, - block_tables, - max_model_len, - ) + Preshuffle = blocksize % 16 == 0 + + if Preshuffle: + kv_num_block, kv_block_Size, _, kv_index_dim = kv_cache_fp8.size() + + split_kv_cache = kv_cache_fp8.view(-1, blocksize * kv_index_dim) + split_kv_cache_data = shuffle_weight( + split_kv_cache[..., : kv_block_Size * index_dim] + .contiguous() + .view([kv_num_block, kv_block_Size, index_dim]) + ) + split_kv_cache[..., : kv_block_Size * index_dim] = ( + split_kv_cache_data.view(kv_num_block, kv_block_Size * index_dim) + ) + _, elapsed_us = run_perftest( deepgemm_fp8_paged_mqa_logits, q_fp8, @@ -286,33 +311,26 @@ def test_deepgemm_fp8_paged_mqa_logits( context_lens, block_tables, max_model_len, - ChunkK, - SplitKV, + ChunkK=ChunkK, + Preshuffle=Preshuffle, + KVBlockSize=blocksize, + WavePerEU=WavePerEU, ) - else: - deepgemm_fp8_paged_mqa_logits_stage1_ragged_k( + cache_key = deepgemm_fp8_paged_mqa_logits( q_fp8, - kv_cache_fp8.view([num_blocks, 1, -1]), - weights, - out_qk, - prefix_sum_context_lens, - kv_indices, - max_model_len, - ) - _, elapsed_us = run_perftest( - deepgemm_fp8_paged_mqa_logits_ragged_k, - q_fp8, - kv_cache_fp8.view([num_blocks, 1, -1]), + kv_cache_fp8, weights, out_logits, - prefix_sum_context_lens, - kv_indices, + context_lens, + block_tables, max_model_len, - ChunkK, - SplitKV, + ChunkK=ChunkK, + Preshuffle=Preshuffle, + KVBlockSize=blocksize, + WavePerEU=WavePerEU, ) - out_qk_logits = torch.sum(out_qk, dim=0) + print(">>> ", cache_key) positions = ( torch.arange(max_model_len, device="cuda") @@ -332,29 +350,40 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor): return 1 - sim out_logits = out_logits.masked_fill(~mask, 0) - out_qk_logits = out_qk_logits.masked_fill(~mask, 0) ref_logits = ref_logits.masked_fill(~mask, 0) - qk_diff = calc_diff(out_qk_logits, ref_logits) logits_diff = calc_diff(out_logits, ref_logits) - assert qk_diff < 1e-3 - assert logits_diff < 1e-3 + print(">>>! logits_diff = ", logits_diff) + # assert logits_diff < 1e-3 total_float_operations = ( 2 * next_n * heads * index_dim * context_lens.float().sum().item() ) flops = total_float_operations / elapsed_us * 1e-6 - ctx_list = context_lens.tolist() - total_memcpyA_bytes = batch_size * next_n * SplitKV * heads * index_dim - total_memcpyB_bytes = ( - sum([cdiv(ctx, ChunkK) * ChunkK * index_dim for ctx in ctx_list]) * next_n + print( + kv_storage_kind, + " time elapsed: ", + elapsed_us, ) - bandwidth_gbps = (total_memcpyA_bytes + total_memcpyB_bytes) / elapsed_us * 1e-3 + if args.aot: + triton_cache_dir = str(triton.knobs.cache.dir) + aot_kernel_dir = f"./paged_mqa_logits/aot" + + padded_str = "T" if args.padding else "F" + os.makedirs(aot_kernel_dir, exist_ok=True) + aot_name = f"paged_mqa_logits{"_preshuffle" if args.kv_preshuffle else ""}_{heads}x{ChunkK}x{index_dim}_B{blocksize}P{padded_str}W{WavePerEU}" - print("bandwidth (GB/s): ", bandwidth_gbps) + src = os.path.join(triton_cache_dir, cache_key) + dst = os.path.join(aot_kernel_dir, aot_name) + if os.path.exists(dst): + os.system(f"rm -rf {dst}") + os.system(f"mv {src} {dst}") + print(f"Moved cache from {src} to {dst}") + + os.system(f"zip -r paged_mqa_logits_aot_kernel paged_mqa_logits") return flops @@ -389,5 +418,32 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor): default=0, help="Q sequence length (mtp + 1 == qo_len) in MTP mode", ) + parser.add_argument( + "-p", + "--padding", + action="store_true", + help="Padding the contiguous dimension of KVCache to multiple of 16 Bytes", + ) + parser.add_argument( + "-aot", + action="store_true", + help="Save compiled triton kernel for later AOT use", + ) + parser.add_argument( + "--perf", + action="store_true", + ) + parser.add_argument( + "--kv_preshuffle", + action="store_true", + help="Enable KV cache preshuffle, also change blocksize to 16", + ) + parser.add_argument( + "--blocksize", + type=int, + default=16, + help="KVCache block size, only used when kv_preshuffle is enabled, must be multiple of 16", + ) + args = parser.parse_args() run_benchmark(args)