Joy-face_change_optimized / deploy /sageattention-cudagraph-stream.patch
kylinfires's picture
Sync main: low-VRAM live test results (624x360)
b9fc164 verified
Raw History Blame Contribute Delete
14.8 kB
diff --git a/csrc/fused/fused.cu b/csrc/fused/fused.cu
index fb8b9f1..00ae804 100644
--- a/csrc/fused/fused.cu
+++ b/csrc/fused/fused.cu
@@ -492,7 +492,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>(),
@@ -574,7 +574,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>(),
@@ -664,7 +664,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>(),
@@ -749,7 +749,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>(),
@@ -834,7 +834,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()),
@@ -911,7 +911,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,
@@ -984,7 +984,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,
@@ -1067,7 +1067,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()),
diff --git a/csrc/qattn/qk_int_sv_f16_cuda_sm80.cu b/csrc/qattn/qk_int_sv_f16_cuda_sm80.cu
index bdedaf6..f3c0347 100644
--- a/csrc/qattn/qk_int_sv_f16_cuda_sm80.cu
+++ b/csrc/qattn/qk_int_sv_f16_cuda_sm80.cu
@@ -14,6 +14,7 @@
* limitations under the License.
*/
+#include <ATen/cuda/CUDAContext.h>
#include "../utils.cuh"
#include <cuda_fp16.h>
#include <cuda_pipeline_primitives.h>
@@ -819,7 +820,7 @@ torch::Tensor qk_int8_sv_f16_accum_f32_attn(torch::Tensor query,
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<half*>(value.data_ptr()),
@@ -994,7 +995,7 @@ torch::Tensor qk_int8_sv_f16_accum_f16_attn(torch::Tensor query,
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<half*>(value.data_ptr()),
@@ -1169,7 +1170,7 @@ torch::Tensor qk_int8_sv_f16_accum_f16_attn_inst_buf(torch::Tensor query,
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<half*>(value.data_ptr()),
@@ -1353,7 +1354,7 @@ torch::Tensor qk_int8_sv_f16_accum_f16_fuse_v_mean_attn(torch::Tensor query,
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<half*>(value.data_ptr()),
diff --git a/csrc/qattn/qk_int_sv_f8_cuda_sm90.cu b/csrc/qattn/qk_int_sv_f8_cuda_sm90.cu
index e9e5ccf..564a1e3 100644
--- a/csrc/qattn/qk_int_sv_f8_cuda_sm90.cu
+++ b/csrc/qattn/qk_int_sv_f8_cuda_sm90.cu
@@ -14,6 +14,7 @@
* limitations under the License.
*/
+#include <ATen/cuda/CUDAContext.h>
#include "../utils.cuh"
#include <cuda.h>
#include <cuda_bf16.h>
@@ -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,
@@ -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,
diff --git a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_attn_inst_buf.cu b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_attn_inst_buf.cu
index c66a81f..2e79322 100644
--- a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_attn_inst_buf.cu
+++ b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_attn_inst_buf.cu
@@ -1,3 +1,4 @@
+#include <ATen/cuda/CUDAContext.h>
#include "attn_cuda_sm89.h"
#include "qk_int_sv_f8_cuda_sm89.cuh"
@@ -152,7 +153,7 @@ torch::Tensor qk_int8_sv_f8_accum_f16_attn_inst_buf(torch::Tensor query,
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<int8_t*>(value.data_ptr()),
diff --git a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu
index 727030f..f19d3e0 100644
--- a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu
+++ b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu
@@ -1,3 +1,4 @@
+#include <ATen/cuda/CUDAContext.h>
#include "attn_cuda_sm89.h"
#include "qk_int_sv_f8_cuda_sm89.cuh"
torch::Tensor qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(torch::Tensor query,
@@ -159,7 +160,7 @@ torch::Tensor qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(torch::Tensor q
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<int8_t*>(value.data_ptr()),
diff --git a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn.cu b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn.cu
index d99d1aa..e3b029c 100644
--- a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn.cu
+++ b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn.cu
@@ -1,3 +1,4 @@
+#include <ATen/cuda/CUDAContext.h>
#include "attn_cuda_sm89.h"
#include "qk_int_sv_f8_cuda_sm89.cuh"
@@ -152,7 +153,7 @@ torch::Tensor qk_int8_sv_f8_accum_f32_attn(torch::Tensor query,
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<int8_t*>(value.data_ptr()),
diff --git a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn_inst_buf.cu b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn_inst_buf.cu
index 14b9b03..12cacb5 100644
--- a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn_inst_buf.cu
+++ b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_attn_inst_buf.cu
@@ -1,3 +1,4 @@
+#include <ATen/cuda/CUDAContext.h>
#include "attn_cuda_sm89.h"
#include "qk_int_sv_f8_cuda_sm89.cuh"
torch::Tensor qk_int8_sv_f8_accum_f32_attn_inst_buf(torch::Tensor query,
@@ -151,7 +152,7 @@ torch::Tensor qk_int8_sv_f8_accum_f32_attn_inst_buf(torch::Tensor query,
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<int8_t*>(value.data_ptr()),
diff --git a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn.cu b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn.cu
index 918b156..035952c 100644
--- a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn.cu
+++ b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn.cu
@@ -1,3 +1,4 @@
+#include <ATen/cuda/CUDAContext.h>
#include "attn_cuda_sm89.h"
#include "qk_int_sv_f8_cuda_sm89.cuh"
torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(torch::Tensor query,
@@ -159,7 +160,7 @@ torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn(torch::Tensor query,
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<int8_t*>(value.data_ptr()),
diff --git a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf.cu b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf.cu
index cc1a85c..6121697 100644
--- a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf.cu
+++ b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf.cu
@@ -1,3 +1,4 @@
+#include <ATen/cuda/CUDAContext.h>
#include "attn_cuda_sm89.h"
#include "qk_int_sv_f8_cuda_sm89.cuh"
torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(torch::Tensor query,
@@ -159,7 +160,7 @@ torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_attn_inst_buf(torch::Tensor q
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<int8_t*>(value.data_ptr()),
diff --git a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn.cu b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn.cu
index fb3db8b..b7f48e2 100644
--- a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn.cu
+++ b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn.cu
@@ -1,3 +1,4 @@
+#include <ATen/cuda/CUDAContext.h>
#include "attn_cuda_sm89.h"
#include "qk_int_sv_f8_cuda_sm89.cuh"
torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(torch::Tensor query,
@@ -164,7 +165,7 @@ torch::Tensor qk_int8_sv_f8_accum_f32_fuse_v_scale_fuse_v_mean_attn(torch::Tenso
dim3 grid(div_ceil(qo_len, CTA_Q), num_qo_heads, batch_size);
dim3 block(32, (CTA_Q / WARP_Q) * (CTA_K / WARP_K));
- kernel_func<<<grid, block, smem_max>>>(
+ kernel_func<<<grid, block, smem_max, at::cuda::getCurrentCUDAStream()>>>(
query.data_ptr<int8_t>(),
key.data_ptr<int8_t>(),
reinterpret_cast<int8_t*>(value.data_ptr()),