@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 148 lines 5.88 KB
Raw
1 // Cooperative cluster TopK for DeepSeek V3 sparse attention indexer.
2 // See cooperative_topk.cuh for kernel implementation.
3
4 #include <cuda_runtime.h>
5
6 #include "torch_utils.h"
7
8 #ifndef USE_ROCM
9 #include "cooperative_topk.cuh"
10 namespace ct = vllm::cooperative;
11 namespace hist4096 = vllm::topk_histogram_4096;
12 #endif
13
14 #ifndef USE_ROCM
15 template <uint32_t TopK, uint32_t CS>
16 void launch_cooperative_cluster(ct::CooperativeTopKParams<TopK>& params,
17 size_t smem, cudaStream_t stream) {
18 auto kernel = []() {
19 if constexpr (CS == 2) {
20 return &ct::cooperative_topk_cs2<TopK>;
21 } else if constexpr (CS == 16) {
22 return &ct::cooperative_topk_cs16<TopK>;
23 } else if constexpr (CS == 8) {
24 return &ct::cooperative_topk_cs8<TopK>;
25 } else {
26 static_assert(CS == 4, "unsupported cooperative_topk cluster size");
27 return &ct::cooperative_topk_cs4<TopK>;
28 }
29 }();
30 if constexpr (CS > 8) {
31 cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed,
32 1);
33 }
34 cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
35 smem);
36
37 cudaLaunchConfig_t cfg = {};
38 cfg.gridDim = dim3(params.num_rows, CS);
39 cfg.blockDim = dim3(hist4096::kBlockSize);
40 cfg.dynamicSmemBytes = smem;
41 cfg.stream = stream;
42 cudaLaunchAttribute attrs[1];
43 attrs[0].id = cudaLaunchAttributeClusterDimension;
44 attrs[0].val.clusterDim = {1, CS, 1};
45 cfg.numAttrs = 1;
46 cfg.attrs = attrs;
47 cudaError_t err = cudaLaunchKernelEx(&cfg, kernel, params);
48 STD_TORCH_CHECK(err == cudaSuccess,
49 "cooperative_topk launch failed: ", cudaGetErrorString(err));
50 }
51
52 template <uint32_t TopK>
53 void launch_cooperative_topk_impl(const torch::stable::Tensor& logits,
54 const torch::stable::Tensor& lengths,
55 torch::stable::Tensor& output,
56 torch::stable::Tensor& workspace,
57 int64_t max_seq_len) {
58 (void)max_seq_len; // Kept for signature parity with persistent_topk.
59 const int64_t num_rows = logits.size(0);
60 const cudaStream_t stream = get_current_cuda_stream();
61
62 const uint32_t stride = static_cast<uint32_t>(logits.stride(0));
63 STD_TORCH_CHECK(
64 num_rows <= 64,
65 "cooperative_topk supports <=64 rows; use persistent_topk for "
66 "larger batches");
67
68 STD_TORCH_CHECK(stride % 4 == 0,
69 "cooperative_topk: stride must be multiple of 4 for TMA "
70 "alignment, got stride (max_model_len)=",
71 stride);
72
73 STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor");
74 STD_TORCH_CHECK(
75 workspace.scalar_type() == torch::headeronly::ScalarType::Byte,
76 "workspace must be uint8");
77
78 ct::CooperativeTopKParams<TopK> params;
79 params.input = logits.const_data_ptr<float>();
80 params.output = output.mutable_data_ptr<int32_t>();
81 params.lengths = lengths.const_data_ptr<int32_t>();
82 params.num_rows = static_cast<uint32_t>(num_rows);
83 params.stride = stride;
84 params.tie_ws =
85 reinterpret_cast<hist4096::Tie*>(workspace.mutable_data_ptr<uint8_t>());
86
87 constexpr uint32_t kTieWsPerRow =
88 TopK <= hist4096::kBlockSize ? hist4096::kMaxTies : TopK;
89 STD_TORCH_CHECK(
90 workspace.size(0) >=
91 static_cast<int64_t>(num_rows * kTieWsPerRow * sizeof(hist4096::Tie)),
92 "workspace too small");
93
94 const bool supports_cluster16 = get_device_prop()->major >= 10;
95 if (num_rows <= 4 && supports_cluster16) {
96 launch_cooperative_cluster<TopK, 16>(params, ct::kSmemSize8, stream);
97 } else if (num_rows <= 8) {
98 launch_cooperative_cluster<TopK, 8>(params, ct::kSmemSize8, stream);
99 } else if (num_rows <= 33) {
100 launch_cooperative_cluster<TopK, 4>(params, ct::kSmemSize4, stream);
101 } else {
102 launch_cooperative_cluster<TopK, 2>(params, ct::kSmemSize2, stream);
103 }
104 }
105 #endif // USE_ROCM
106
107 void cooperative_topk(const torch::stable::Tensor& logits,
108 const torch::stable::Tensor& lengths,
109 torch::stable::Tensor& output,
110 torch::stable::Tensor& workspace, int64_t k,
111 int64_t max_seq_len) {
112 #ifndef USE_ROCM
113 STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor");
114 STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor");
115 STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor");
116 STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float,
117 "Only float32 supported");
118 STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int,
119 "lengths must be int32");
120 STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int,
121 "output must be int32");
122 STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D");
123 STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2,
124 "lengths must be 1D or 2D");
125 STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous");
126 STD_TORCH_CHECK(output.dim() == 2, "output must be 2D");
127 const int64_t num_rows = logits.size(0);
128 STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch");
129 STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k,
130 "output size mismatch");
131 STD_TORCH_CHECK(
132 k == 512 || k == 1024 || k == 2048,
133 "cooperative_topk supports k=512, k=1024, or k=2048, got k=", k);
134
135 if (k == 512) {
136 launch_cooperative_topk_impl<512>(logits, lengths, output, workspace,
137 max_seq_len);
138 } else if (k == 1024) {
139 launch_cooperative_topk_impl<1024>(logits, lengths, output, workspace,
140 max_seq_len);
141 } else {
142 launch_cooperative_topk_impl<2048>(logits, lengths, output, workspace,
143 max_seq_len);
144 }
145 #else
146 STD_TORCH_CHECK(false, "cooperative_topk is not supported on ROCm");
147 #endif
148 }