Skip to content

Launch sm90 attention and fused quant kernels on the current torch stream (CUDA-graph capture correctness) #371

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shanai13 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Launch sm90 attention and fused quant kernels on the current torch stream (CUDA-graph capture correctness) #371

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 9 additions & 8 deletions csrc/fused/fused.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

#include <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#include "../dispatch_utils.h"
#include "../utils.cuh"
Expand Down Expand Up @@ -492,7 +493,7 @@ void quant_per_block_int8_cuda(

dim3 block(BLOCK_SIZE * (HEAD_DIM / 8) / num_pack_per_thread);

QuantInt8Kernel<HEAD_DIM, BLOCK_SIZE, num_pack_per_thread, true, false, c_type><<<grid, block>>>(
QuantInt8Kernel<HEAD_DIM, BLOCK_SIZE, num_pack_per_thread, true, false, c_type><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
reinterpret_cast<c_type*>(input.data_ptr()),
nullptr,
output.data_ptr<int8_t>(),
Expand Down Expand Up @@ -574,7 +575,7 @@ void quant_per_block_int8_cuda(

dim3 block(BLOCK_SIZE * (HEAD_DIM / 8) / num_pack_per_thread);

QuantInt8Kernel<HEAD_DIM, BLOCK_SIZE, num_pack_per_thread, false, false, c_type><<<grid, block>>>(
QuantInt8Kernel<HEAD_DIM, BLOCK_SIZE, num_pack_per_thread, false, false, c_type><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
reinterpret_cast<c_type*>(input.data_ptr()),
nullptr,
output.data_ptr<int8_t>(),
Expand Down Expand Up @@ -664,7 +665,7 @@ void quant_per_block_int8_fuse_sub_mean_cuda(

dim3 block(BLOCK_SIZE * (HEAD_DIM / 8) / num_pack_per_thread);

QuantInt8Kernel<HEAD_DIM, BLOCK_SIZE, num_pack_per_thread, false, true, c_type><<<grid, block>>>(
QuantInt8Kernel<HEAD_DIM, BLOCK_SIZE, num_pack_per_thread, false, true, c_type><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
reinterpret_cast<c_type*>(input.data_ptr()),
reinterpret_cast<c_type*>(mean.data_ptr()),
output.data_ptr<int8_t>(),
Expand Down Expand Up @@ -749,7 +750,7 @@ void quant_per_warp_int8_cuda(

dim3 block(WARP_BLOCK_SIZE * (HEAD_DIM / 8) / num_pack_per_thread);

QuantInt8Kernel<HEAD_DIM, WARP_BLOCK_SIZE, num_pack_per_thread, false, false, c_type><<<grid, block>>>(
QuantInt8Kernel<HEAD_DIM, WARP_BLOCK_SIZE, num_pack_per_thread, false, false, c_type><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
reinterpret_cast<c_type*>(input.data_ptr()),
nullptr,
output.data_ptr<int8_t>(),
Expand Down Expand Up @@ -834,7 +835,7 @@ void sub_mean_cuda(

dim3 block(BLOCK_SIZE * (HEAD_DIM / 8) / num_pack_per_thread);

SubMeanKernel<HEAD_DIM, BLOCK_SIZE, num_pack_per_thread><<<grid, block>>>(
SubMeanKernel<HEAD_DIM, BLOCK_SIZE, num_pack_per_thread><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
reinterpret_cast<c_type*>(input.data_ptr()),
reinterpret_cast<c_type*>(mean.data_ptr()),
reinterpret_cast<half*>(output.data_ptr()),
Expand Down Expand Up @@ -911,7 +912,7 @@ void transpose_pad_permute_cuda(

dim3 block(CTA_SIZE * (HEAD_DIM / 8));

TransposePadPermuteKernel<HEAD_DIM, CTA_SIZE, true, c_type><<<grid, block>>>(
TransposePadPermuteKernel<HEAD_DIM, CTA_SIZE, true, c_type><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
reinterpret_cast<c_type*>(input.data_ptr()),
reinterpret_cast<c_type*>(output.data_ptr()),
num_tokens,
Expand Down Expand Up @@ -984,7 +985,7 @@ void scale_fuse_quant_cuda(
auto input_dtype = input.scalar_type();

DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input_dtype, c_type, {
MeanScaleKernel<64, false, c_type><<<grid, block>>>(
MeanScaleKernel<64, false, c_type><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
reinterpret_cast<c_type*>(input.data_ptr()),
reinterpret_cast<int8_t*>(output.data_ptr()),
nullptr,
Expand Down Expand Up @@ -1067,7 +1068,7 @@ void mean_scale_fuse_quant_cuda(
auto input_dtype = input.scalar_type();

DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FP16(input_dtype, c_type, {
MeanScaleKernel<64, true, c_type><<<grid, block>>>(
MeanScaleKernel<64, true, c_type><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
reinterpret_cast<c_type*>(input.data_ptr()),
reinterpret_cast<int8_t*>(output.data_ptr()),
reinterpret_cast<float*>(mean.data_ptr()),
Expand Down
5 changes: 3 additions & 2 deletions csrc/qattn/qk_int_sv_f8_cuda_sm90.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
#include <cuda_bf16.h>
#include <cuda_fp8.h>
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>

#include "../wgmma.cuh"
#include "../math.cuh"
Expand Down Expand Up @@ -717,7 +718,7 @@ torch::Tensor qk_int8_sv_f8_accum_f32_attn_inst_buf(
cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize);

dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
kernel<<<grid, NUM_THREADS, sMemSize>>>(
kernel<<<grid, NUM_THREADS, sMemSize, at::cuda::getCurrentCUDAStream()>>>(
tma_map_Q,
tma_map_K,
tma_map_V,
Expand Down Expand Up @@ -895,7 +896,7 @@ torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(
cudaFuncAttributeMaxDynamicSharedMemorySize, sMemSize);

dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
kernel<<<grid, NUM_THREADS, sMemSize>>>(
kernel<<<grid, NUM_THREADS, sMemSize, at::cuda::getCurrentCUDAStream()>>>(
tma_map_Q,
tma_map_K,
tma_map_V,
Expand Down