@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 822 lines 41.5 KB
Raw
1 #pragma once
2
3 #include <torch/csrc/stable/library.h>
4 #include <torch/csrc/stable/tensor.h>
5 #include <torch/headeronly/util/Exception.h>
6
7 #include <optional>
8 #include <string>
9 #include <vector>
10
11 #include <torch/csrc/stable/ops.h>
12
13 inline torch::stable::Tensor weak_ref_tensor(torch::stable::Tensor& tensor) {
14 // Ensure tensor is on CUDA
15 STD_TORCH_CHECK(tensor.device().is_cuda(), "Tensor must be on CUDA device");
16
17 // Get the raw data pointer
18 void* data_ptr = tensor.mutable_data_ptr();
19
20 /// Create a new tensor from the raw data pointer
21 return torch::stable::from_blob(data_ptr, tensor.sizes(), tensor.strides(),
22 tensor.device(), tensor.scalar_type());
23 }
24
25 void per_token_group_quant_fp8(const torch::stable::Tensor& input,
26 torch::stable::Tensor& output_q,
27 torch::stable::Tensor& output_s,
28 int64_t group_size, double eps, double fp8_min,
29 double fp8_max, bool scale_ue8m0,
30 bool dummy_is_scale_transposed,
31 bool dummy_is_tma_aligned);
32
33 // Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales.
34 void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input,
35 torch::stable::Tensor& output_q,
36 torch::stable::Tensor& output_s_packed,
37 int64_t group_size, double eps,
38 double min_8bit, double max_8bit);
39
40 void per_token_group_quant_int8(const torch::stable::Tensor& input,
41 torch::stable::Tensor& output_q,
42 torch::stable::Tensor& output_s,
43 int64_t group_size, double eps, double int8_min,
44 double int8_max);
45
46 #ifndef USE_ROCM
47 bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
48 bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
49 bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
50
51 void cutlass_scaled_mm(torch::stable::Tensor& out,
52 torch::stable::Tensor const& a,
53 torch::stable::Tensor const& b,
54 torch::stable::Tensor const& a_scales,
55 torch::stable::Tensor const& b_scales,
56 std::optional<torch::stable::Tensor> const& bias);
57
58 void cutlass_moe_mm(torch::stable::Tensor& out_tensors,
59 torch::stable::Tensor const& a_tensors,
60 torch::stable::Tensor const& b_tensors,
61 torch::stable::Tensor const& a_scales,
62 torch::stable::Tensor const& b_scales,
63 torch::stable::Tensor const& expert_offsets,
64 torch::stable::Tensor const& problem_sizes,
65 torch::stable::Tensor const& a_strides,
66 torch::stable::Tensor const& b_strides,
67 torch::stable::Tensor const& c_strides, bool per_act_token,
68 bool per_out_ch);
69
70 void cutlass_scaled_mm_azp(torch::stable::Tensor& out,
71 torch::stable::Tensor const& a,
72 torch::stable::Tensor const& b,
73 torch::stable::Tensor const& a_scales,
74 torch::stable::Tensor const& b_scales,
75 torch::stable::Tensor const& azp_adj,
76 std::optional<torch::stable::Tensor> const& azp,
77 std::optional<torch::stable::Tensor> const& bias);
78
79 void get_cutlass_moe_mm_data(
80 const torch::stable::Tensor& topk_ids,
81 torch::stable::Tensor& expert_offsets,
82 torch::stable::Tensor& problem_sizes1,
83 torch::stable::Tensor& problem_sizes2,
84 torch::stable::Tensor& input_permutation,
85 torch::stable::Tensor& output_permutation, const int64_t num_experts,
86 const int64_t n, const int64_t k,
87 const std::optional<torch::stable::Tensor>& blockscale_offsets,
88 const bool is_gated);
89
90 void get_cutlass_moe_mm_problem_sizes_from_expert_offsets(
91 const torch::stable::Tensor& expert_first_token_offset,
92 torch::stable::Tensor& problem_sizes1,
93 torch::stable::Tensor& problem_sizes2, const int64_t n, const int64_t k,
94 const bool swap_ab);
95
96 void get_cutlass_batched_moe_mm_data(
97 torch::stable::Tensor& expert_offsets,
98 torch::stable::Tensor& problem_sizes1,
99 torch::stable::Tensor& problem_sizes2,
100 const torch::stable::Tensor& expert_num_tokens,
101 const int64_t num_local_experts, const int64_t padded_m, const int64_t n,
102 const int64_t k);
103
104 // FP4/NVFP4 ops
105 bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability);
106
107 void cutlass_scaled_fp4_mm(torch::stable::Tensor& D,
108 torch::stable::Tensor const& A,
109 torch::stable::Tensor const& B,
110 torch::stable::Tensor const& A_sf,
111 torch::stable::Tensor const& B_sf,
112 torch::stable::Tensor const& alpha);
113
114 void cutlass_fp4_group_mm(torch::stable::Tensor& output,
115 const torch::stable::Tensor& a,
116 const torch::stable::Tensor& b,
117 const torch::stable::Tensor& a_blockscale,
118 const torch::stable::Tensor& b_blockscales,
119 const torch::stable::Tensor& alphas,
120 const torch::stable::Tensor& problem_sizes,
121 const torch::stable::Tensor& expert_offsets,
122 const torch::stable::Tensor& sf_offsets);
123
124 std::tuple<torch::stable::Tensor, torch::stable::Tensor> scaled_fp4_quant_func(
125 torch::stable::Tensor const& input,
126 torch::stable::Tensor const& input_scale, bool is_sf_swizzled_layout);
127
128 void scaled_fp4_quant_out(torch::stable::Tensor const& input,
129 torch::stable::Tensor const& input_scale,
130 bool is_sf_swizzled_layout,
131 torch::stable::Tensor& output,
132 torch::stable::Tensor& output_scale);
133
134 void scaled_fp4_experts_quant(
135 torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
136 torch::stable::Tensor const& input,
137 torch::stable::Tensor const& input_global_scale,
138 torch::stable::Tensor const& input_offset_by_experts,
139 torch::stable::Tensor const& output_scale_offset_by_experts);
140
141 void silu_and_mul_scaled_fp4_experts_quant(
142 torch::stable::Tensor& output, torch::stable::Tensor& output_scale,
143 torch::stable::Tensor const& input,
144 torch::stable::Tensor const& input_global_scale,
145 torch::stable::Tensor const& input_offset_by_experts,
146 torch::stable::Tensor const& output_scale_offset_by_experts);
147
148 void silu_and_mul_nvfp4_quant(torch::stable::Tensor& out,
149 torch::stable::Tensor& output_block_scale,
150 torch::stable::Tensor& input,
151 torch::stable::Tensor& input_global_scale);
152
153 void cutlass_mxfp4_group_mm(torch::stable::Tensor& output,
154 const torch::stable::Tensor& a,
155 const torch::stable::Tensor& b,
156 const torch::stable::Tensor& a_blockscale,
157 const torch::stable::Tensor& b_blockscales,
158 const torch::stable::Tensor& problem_sizes,
159 const torch::stable::Tensor& expert_offsets,
160 const torch::stable::Tensor& sf_offsets);
161
162 // AWQ ops
163 torch::stable::Tensor awq_gemm(torch::stable::Tensor _in_feats,
164 torch::stable::Tensor _kernel,
165 torch::stable::Tensor _scaling_factors,
166 torch::stable::Tensor _zeros,
167 int64_t split_k_iters);
168
169 torch::stable::Tensor awq_dequantize(torch::stable::Tensor _kernel,
170 torch::stable::Tensor _scaling_factors,
171 torch::stable::Tensor _zeros,
172 int64_t split_k_iters, int64_t thx,
173 int64_t thy);
174
175 // DSV3 fused A GEMM: conditionally compiled so declaration and impl
176 // registration are in the source file (dsv3_fused_a_gemm.cu)
177
178 #endif
179
180 // CPU tensor -> CUDA UVA view (shared CUDA/ROCm)
181 torch::stable::Tensor get_cuda_view_from_cpu_tensor(
182 torch::stable::Tensor& cpu_tensor);
183
184 // Attention kernels (shared CUDA/ROCm)
185 void merge_attn_states(
186 torch::stable::Tensor& output,
187 std::optional<torch::stable::Tensor> output_lse,
188 const torch::stable::Tensor& prefix_output,
189 const torch::stable::Tensor& prefix_lse,
190 const torch::stable::Tensor& suffix_output,
191 const torch::stable::Tensor& suffix_lse,
192 const std::optional<int64_t> prefill_tokens_with_context,
193 const std::optional<torch::stable::Tensor>& output_scale = std::nullopt);
194
195 torch::stable::Tensor hadacore_transform(torch::stable::Tensor& x,
196 bool inplace);
197
198 // Layernorm kernels (shared CUDA/ROCm)
199 void rms_norm(torch::stable::Tensor& out, torch::stable::Tensor& input,
200 std::optional<torch::stable::Tensor> weight, double epsilon);
201
202 void fused_add_rms_norm(torch::stable::Tensor& input,
203 torch::stable::Tensor& residual,
204 std::optional<torch::stable::Tensor> weight,
205 double epsilon);
206
207 // Layernorm-quant kernels (shared CUDA/ROCm)
208 void rms_norm_static_fp8_quant(torch::stable::Tensor& out,
209 torch::stable::Tensor& input,
210 torch::stable::Tensor& weight,
211 torch::stable::Tensor& scale, double epsilon);
212
213 void fused_add_rms_norm_static_fp8_quant(torch::stable::Tensor& out,
214 torch::stable::Tensor& input,
215 torch::stable::Tensor& residual,
216 torch::stable::Tensor& weight,
217 torch::stable::Tensor& scale,
218 double epsilon);
219
220 // Fused layernorm + dynamic per-token quant kernels (shared CUDA/ROCm)
221 void rms_norm_dynamic_per_token_quant(
222 torch::stable::Tensor& out, torch::stable::Tensor const& input,
223 torch::stable::Tensor const& weight, torch::stable::Tensor& scales,
224 double const var_epsilon, std::optional<torch::stable::Tensor> scale_ub,
225 std::optional<torch::stable::Tensor> residual);
226
227 void rms_norm_per_block_quant(torch::stable::Tensor& out,
228 torch::stable::Tensor const& input,
229 torch::stable::Tensor const& weight,
230 torch::stable::Tensor& scales,
231 double const var_epsilon,
232 std::optional<torch::stable::Tensor> scale_ub,
233 std::optional<torch::stable::Tensor> residual,
234 int64_t group_size, bool is_scale_transposed);
235
236 void silu_and_mul_per_block_quant(torch::stable::Tensor& out,
237 torch::stable::Tensor const& input,
238 torch::stable::Tensor& scales,
239 int64_t group_size,
240 std::optional<torch::stable::Tensor> scale_ub,
241 bool is_scale_transposed);
242
243 // Positional encoding kernels (shared CUDA/ROCm)
244 void rotary_embedding(torch::stable::Tensor& positions,
245 torch::stable::Tensor& query,
246 std::optional<torch::stable::Tensor> key,
247 int64_t head_size, torch::stable::Tensor& cos_sin_cache,
248 bool is_neox, int64_t rope_dim_offset, bool inverse);
249
250 void fused_qk_norm_rope(torch::stable::Tensor& qkv, int64_t num_heads_q,
251 int64_t num_heads_k, int64_t num_heads_v,
252 int64_t head_dim, double eps,
253 torch::stable::Tensor& q_weight,
254 torch::stable::Tensor& k_weight,
255 torch::stable::Tensor& cos_sin_cache, bool is_neox,
256 torch::stable::Tensor& position_ids,
257 int64_t forced_token_heads_per_warp);
258
259 void fused_deepseek_v4_kv_rope_insert(
260 torch::stable::Tensor const& kv, torch::stable::Tensor& k_cache,
261 torch::stable::Tensor const& slot_mapping,
262 torch::stable::Tensor const& position_ids,
263 torch::stable::Tensor const& cos_sin_cache, int64_t cache_block_size,
264 std::optional<torch::stable::Tensor> fp8_scale, bool kv_mxfp8);
265
266 torch::stable::Tensor fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
267 torch::stable::Tensor const& q_in, torch::stable::Tensor const& kv,
268 torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
269 torch::stable::Tensor const& position_ids,
270 torch::stable::Tensor const& cos_sin_cache, int64_t q_head_padded,
271 double eps, int64_t cache_block_size, bool apply_q_norm, bool kv_mxfp8,
272 bool apply_q_rope, bool is_q_interleaved);
273
274 void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert(
275 torch::stable::Tensor& q, torch::stable::Tensor const& kv,
276 torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
277 torch::stable::Tensor const& position_ids,
278 torch::stable::Tensor const& cos_sin_cache, double eps,
279 int64_t cache_block_size, bool apply_q_norm);
280
281 void fused_kimi_k3_mla_key_concat_kv_cache_insert(
282 torch::stable::Tensor& q, torch::stable::Tensor const& k_nope,
283 torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
284 torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache,
285 torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
286 std::optional<torch::stable::Tensor> position_ids,
287 std::optional<torch::stable::Tensor> cos_sin_cache);
288
289 void fused_kimi_k3_mla_key_concat_ds_mla_insert(
290 torch::stable::Tensor& q, torch::stable::Tensor const& k_nope,
291 torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
292 torch::stable::Tensor& k_out, torch::stable::Tensor& k_cache,
293 torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
294 std::optional<torch::stable::Tensor> position_ids,
295 std::optional<torch::stable::Tensor> cos_sin_cache);
296
297 void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope,
298 torch::stable::Tensor const& k_pe,
299 torch::stable::Tensor& k_out);
300
301 void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope,
302 torch::stable::Tensor const& k_pe,
303 torch::stable::Tensor const& v,
304 torch::stable::Tensor& k_fp8,
305 torch::stable::Tensor& v_fp8);
306
307 void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(
308 torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope,
309 torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed,
310 torch::stable::Tensor const& v, torch::stable::Tensor& q_fp8,
311 torch::stable::Tensor& k_fp8, torch::stable::Tensor& v_fp8,
312 torch::stable::Tensor& k_cache, torch::stable::Tensor const& slot_mapping,
313 torch::stable::Tensor const& q_scale_inv,
314 torch::stable::Tensor const& k_scale_inv,
315 torch::stable::Tensor const& v_scale_inv,
316 torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size,
317 std::optional<torch::stable::Tensor> position_ids,
318 std::optional<torch::stable::Tensor> cos_sin_cache);
319
320 void fused_kimi_k3_mla_decode_q_concat_kv_cache_insert(
321 torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
322 torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
323 torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
324 torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
325 std::optional<torch::stable::Tensor> position_ids,
326 std::optional<torch::stable::Tensor> cos_sin_cache);
327
328 void fused_kimi_k3_mla_decode_q_concat_kv_cache_fp8_insert(
329 torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
330 torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
331 torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
332 torch::stable::Tensor const& slot_mapping,
333 torch::stable::Tensor const& q_scale_inv,
334 torch::stable::Tensor const& cache_scale_inv, int64_t cache_block_size,
335 std::optional<torch::stable::Tensor> position_ids,
336 std::optional<torch::stable::Tensor> cos_sin_cache);
337
338 void fused_kimi_k3_mla_decode_q_concat_ds_mla_insert(
339 torch::stable::Tensor const& ql_nope, torch::stable::Tensor const& q_pe,
340 torch::stable::Tensor const& kv_c_normed, torch::stable::Tensor const& k_pe,
341 torch::stable::Tensor& mqa_q, torch::stable::Tensor& k_cache,
342 torch::stable::Tensor const& slot_mapping, int64_t cache_block_size,
343 std::optional<torch::stable::Tensor> position_ids,
344 std::optional<torch::stable::Tensor> cos_sin_cache);
345
346 void fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert(
347 torch::stable::Tensor const& q, torch::stable::Tensor const& kv,
348 torch::stable::Tensor& q_fp8, torch::stable::Tensor& k_cache,
349 torch::stable::Tensor const& slot_mapping,
350 torch::stable::Tensor const& position_ids,
351 torch::stable::Tensor const& cos_sin_cache,
352 torch::stable::Tensor const& fp8_scale,
353 torch::stable::Tensor const& q_fp8_scale_inv, double eps,
354 int64_t cache_block_size, bool apply_q_norm);
355
356 #ifndef USE_ROCM
357 std::tuple<torch::stable::Tensor, torch::stable::Tensor>
358 minimax_allreduce_rms_qk(torch::stable::Tensor qkv,
359 torch::stable::Tensor const& norm_weight_q,
360 torch::stable::Tensor const& norm_weight_k,
361 torch::stable::Tensor workspace, int64_t const q_size,
362 int64_t const kv_size, int64_t const rank,
363 int64_t const nranks, double const eps);
364 #endif
365
366 // Horizontally-fused MiniMax-M3 QK-norm + partial NeoX RoPE (+ optional KV /
367 // index-cache insert). Dense layer: norm+RoPE only; sparse layer: also packs
368 // the index branch and scatters k/v/index_k into their paged caches.
369 void fused_minimax_m3_qknorm_rope_kv_insert(
370 torch::stable::Tensor& qkv, torch::stable::Tensor const& q_norm_weight,
371 torch::stable::Tensor const& k_norm_weight,
372 torch::stable::Tensor const& cos_sin_cache,
373 torch::stable::Tensor const& positions, int64_t num_heads,
374 int64_t num_kv_heads, int64_t rotary_dim, double eps,
375 std::optional<torch::stable::Tensor> index_q_norm_weight,
376 std::optional<torch::stable::Tensor> index_k_norm_weight,
377 int64_t num_index_heads, std::optional<torch::stable::Tensor> slot_mapping,
378 std::optional<torch::stable::Tensor> index_slot_mapping,
379 std::optional<torch::stable::Tensor> kv_cache,
380 std::optional<torch::stable::Tensor> index_cache, int64_t block_size,
381 std::optional<torch::stable::Tensor> q_out,
382 std::optional<torch::stable::Tensor> index_q_out,
383 const std::string& kv_cache_dtype, bool skip_index_branch,
384 std::optional<torch::stable::Tensor> q_fp8_out, double q_fp8_scale,
385 std::optional<torch::stable::Tensor> kv_k_scale,
386 std::optional<torch::stable::Tensor> kv_v_scale);
387
388 #ifdef VLLM_ENABLE_FUSED_KDA_DECODE
389 void fused_kda_decode(
390 torch::stable::Tensor const& x, torch::stable::Tensor const& weight,
391 std::optional<torch::stable::Tensor> bias,
392 torch::stable::Tensor& conv_state, torch::stable::Tensor const& raw_g,
393 torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log,
394 torch::stable::Tensor const& dt_bias,
395 torch::stable::Tensor const& state_indices, torch::stable::Tensor& state,
396 torch::stable::Tensor& out, std::optional<double> lower_bound,
397 std::optional<torch::stable::Tensor> output_gate,
398 std::optional<torch::stable::Tensor> norm_weight, double norm_eps);
399
400 #endif
401
402 #ifdef VLLM_ENABLE_FUSED_GDN_DECODE
403 void fused_gdn_decode_post_conv_mtp(
404 torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a,
405 torch::stable::Tensor const& b, torch::stable::Tensor const& a_log,
406 torch::stable::Tensor const& dt_bias,
407 torch::stable::Tensor const& state_indices,
408 torch::stable::Tensor const& cu_seqlens,
409 torch::stable::Tensor const& num_accepted_tokens,
410 torch::stable::Tensor& state, torch::stable::Tensor const& output_gate,
411 torch::stable::Tensor const& norm_weight, torch::stable::Tensor& out,
412 double scale, double norm_eps, const std::string& output_gate_activation);
413
414 #endif
415
416 #ifdef VLLM_ENABLE_FUSED_KDA_CHUNK
417 void fused_kda_prologue(
418 torch::stable::Tensor const& q, torch::stable::Tensor const& k,
419 torch::stable::Tensor const& v, torch::stable::Tensor const& raw_g,
420 torch::stable::Tensor const& raw_beta, torch::stable::Tensor const& a_log,
421 torch::stable::Tensor const& dt_bias, torch::stable::Tensor& qg,
422 torch::stable::Tensor& w, torch::stable::Tensor& u,
423 torch::stable::Tensor& kg_t, torch::stable::Tensor& aqk,
424 torch::stable::Tensor& decay, torch::stable::Tensor const& cu_seqlens,
425 torch::stable::Tensor const& chunk_indices,
426 std::optional<torch::stable::Tensor> conv_weight,
427 std::optional<torch::stable::Tensor> conv_state,
428 std::optional<torch::stable::Tensor> conv_state_indices,
429 std::optional<torch::stable::Tensor> conv_has_initial_state, double scale,
430 double lower_bound);
431
432 void fused_kda_chunk(
433 torch::stable::Tensor const& qg, torch::stable::Tensor const& w,
434 torch::stable::Tensor const& u, torch::stable::Tensor const& kg_t,
435 torch::stable::Tensor const& aqk, torch::stable::Tensor const& decay,
436 std::optional<torch::stable::Tensor> initial_state,
437 std::optional<torch::stable::Tensor> final_state,
438 torch::stable::Tensor& out, torch::stable::Tensor const& cu_seqlens,
439 torch::stable::Tensor const& chunk_offsets, double scale,
440 std::optional<torch::stable::Tensor> group_state, int64_t groups,
441 std::optional<torch::stable::Tensor> checkpoint_state,
442 std::optional<torch::stable::Tensor> checkpoint_offsets,
443 std::optional<torch::stable::Tensor> checkpoint_state_indices,
444 std::optional<torch::stable::Tensor> state_cache,
445 std::optional<torch::stable::Tensor> state_indices,
446 std::optional<torch::stable::Tensor> has_initial_state);
447 #endif
448
449 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
450 void kimi_k3_attn_res(torch::stable::Tensor& prefix,
451 std::optional<torch::stable::Tensor> delta,
452 torch::stable::Tensor& blocks,
453 torch::stable::Tensor const& norm_weight,
454 torch::stable::Tensor const& qk_weight,
455 std::optional<torch::stable::Tensor> output_norm_weight,
456 torch::stable::Tensor& output, int64_t num_blocks,
457 int64_t block_write_idx, double eps,
458 double output_norm_eps);
459 #endif
460
461 // Sampler kernels (shared CUDA/ROCm)
462 void apply_repetition_penalties_(
463 torch::stable::Tensor& logits, const torch::stable::Tensor& prompt_mask,
464 const torch::stable::Tensor& output_mask,
465 const torch::stable::Tensor& repetition_penalties);
466
467 void top_k_per_row_prefill(const torch::stable::Tensor& logits,
468 const torch::stable::Tensor& rowStarts,
469 const torch::stable::Tensor& rowEnds,
470 torch::stable::Tensor& indices, int64_t numRows,
471 int64_t stride0, int64_t stride1, int64_t topK);
472
473 void top_k_per_row_decode(const torch::stable::Tensor& logits, int64_t next_n,
474 const torch::stable::Tensor& seqLens,
475 torch::stable::Tensor& indices, int64_t numRows,
476 int64_t stride0, int64_t stride1, int64_t topK);
477
478 void persistent_topk(const torch::stable::Tensor& logits,
479 const torch::stable::Tensor& lengths,
480 torch::stable::Tensor& output,
481 torch::stable::Tensor& workspace, int64_t k,
482 int64_t max_seq_len);
483
484 #ifdef VLLM_ENABLE_COOPERATIVE_TOPK
485 void cooperative_topk(const torch::stable::Tensor& logits,
486 const torch::stable::Tensor& lengths,
487 torch::stable::Tensor& output,
488 torch::stable::Tensor& workspace, int64_t k,
489 int64_t max_seq_len);
490 #endif
491
492 void selective_scan_fwd(
493 const torch::stable::Tensor& u, const torch::stable::Tensor& delta,
494 const torch::stable::Tensor& A, const torch::stable::Tensor& B,
495 const torch::stable::Tensor& C,
496 const std::optional<torch::stable::Tensor>& D_,
497 const std::optional<torch::stable::Tensor>& z_,
498 const std::optional<torch::stable::Tensor>& delta_bias_,
499 bool delta_softplus,
500 const std::optional<torch::stable::Tensor>& query_start_loc,
501 const std::optional<torch::stable::Tensor>& cache_indices,
502 const std::optional<torch::stable::Tensor>& has_initial_state,
503 const torch::stable::Tensor& ssm_states, int64_t null_block_id,
504 int64_t block_size,
505 const std::optional<torch::stable::Tensor>& block_idx_first_scheduled_token,
506 const std::optional<torch::stable::Tensor>& block_idx_last_scheduled_token,
507 const std::optional<torch::stable::Tensor>& initial_state_idx,
508 const std::optional<torch::stable::Tensor>& cu_chunk_seqlen,
509 const std::optional<torch::stable::Tensor>& last_chunk_indices);
510
511 using fptr_t = int64_t;
512 fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
513 torch::stable::Tensor& rank_data, int64_t rank,
514 bool fully_connected);
515 void all_reduce(fptr_t _fa, torch::stable::Tensor& inp,
516 torch::stable::Tensor& out, fptr_t reg_buffer,
517 int64_t reg_buffer_sz_bytes);
518 void custom_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
519 torch::stable::Tensor& out, fptr_t reg_buffer,
520 int64_t reg_buffer_sz_bytes);
521 void mnnvl_lamport_all_gather(fptr_t _fa, torch::stable::Tensor& inp,
522 torch::stable::Tensor& out, fptr_t local_buffer,
523 fptr_t multicast_buffer, fptr_t epoch_buffer,
524 int64_t stage_sz_bytes);
525 void custom_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
526 torch::stable::Tensor& out, fptr_t reg_buffer,
527 int64_t reg_buffer_sz_bytes);
528 void mnnvl_lamport_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
529 torch::stable::Tensor& out,
530 fptr_t local_buffer, fptr_t epoch_buffer,
531 int64_t stage_sz_bytes);
532 void mnnvl_multimem_reduce_scatter(fptr_t _fa, torch::stable::Tensor& inp,
533 torch::stable::Tensor& out,
534 fptr_t local_buffer, fptr_t multicast_buffer,
535 int64_t stage_sz_bytes, int64_t block_limit);
536 void dispose(fptr_t _fa);
537 int64_t meta_size();
538 void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);
539 std::tuple<std::vector<int64_t>, std::vector<int64_t>>
540 get_graph_buffer_ipc_meta(fptr_t _fa);
541 void register_graph_buffers(fptr_t _fa,
542 const std::vector<std::vector<int64_t>>& handles,
543 const std::vector<std::vector<int64_t>>& offsets);
544 std::tuple<int64_t, torch::stable::Tensor> allocate_shared_buffer_and_handle(
545 int64_t size);
546 int64_t open_mem_handle(torch::stable::Tensor& mem_handle);
547 void free_shared_buffer(int64_t buffer);
548
549 // Activation kernels (shared CUDA/ROCm)
550 void silu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
551 void silu_and_mul_clamp(torch::stable::Tensor& out,
552 torch::stable::Tensor& input, double limit,
553 double alpha = 1.0, double beta = 0.0);
554
555 void silu_and_mul_quant(torch::stable::Tensor& out,
556 torch::stable::Tensor& input,
557 torch::stable::Tensor& scale);
558
559 void persistent_masked_m_silu_mul_quant(
560 const torch::stable::Tensor& input, // (E, T, 2*H)
561 const torch::stable::Tensor& tokens_per_expert, // (E)
562 torch::stable::Tensor& y_q, // (E, T, H) [OUT]
563 torch::stable::Tensor& y_s, // (E, T, H//group_size) [OUT]
564 bool use_ue8m0);
565
566 void mul_and_silu(torch::stable::Tensor& out, torch::stable::Tensor& input);
567 void gelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input);
568 void gelu_tanh_and_mul(torch::stable::Tensor& out,
569 torch::stable::Tensor& input);
570 void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
571 double threshold);
572 void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
573 double alpha = 1.702, double limit = 7.0);
574 void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
575 double beta = 1.0, double linear_beta = -1.0);
576 void situ_and_mul_quant(
577 torch::stable::Tensor& out, torch::stable::Tensor& scale,
578 torch::stable::Tensor& input, double beta = 1.0, double linear_beta = -1.0,
579 int64_t group_size = 0,
580 std::optional<torch::stable::Tensor> num_valid_tokens = std::nullopt,
581 int64_t topk = 1);
582 void masked_situ_and_mul(torch::stable::Tensor& out,
583 torch::stable::Tensor& input,
584 const torch::stable::Tensor& expert_num_tokens,
585 double beta = 1.0, double linear_beta = -1.0);
586 void masked_moe_activation(torch::stable::Tensor& out,
587 torch::stable::Tensor& input,
588 const torch::stable::Tensor& valid_token_counts,
589 const std::string& activation,
590 double clamp_limit = 0.0, double alpha = 1.0,
591 double beta = 0.0, double situ_beta = 1.0,
592 double situ_linear_beta = -1.0);
593 void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
594 void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
595 void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
596
597 void relu_squared(torch::stable::Tensor& out, torch::stable::Tensor& input);
598
599 // INT8 quantization kernels (shared CUDA/ROCm)
600 void static_scaled_int8_quant(torch::stable::Tensor& out,
601 torch::stable::Tensor const& input,
602 torch::stable::Tensor const& scale,
603 std::optional<torch::stable::Tensor> const& azp);
604
605 void dynamic_scaled_int8_quant(torch::stable::Tensor& out,
606 torch::stable::Tensor const& input,
607 torch::stable::Tensor& scales,
608 std::optional<torch::stable::Tensor> const& azp);
609
610 // FP8 quantization kernels (shared CUDA/ROCm)
611 void static_scaled_fp8_quant(
612 torch::stable::Tensor& out, torch::stable::Tensor const& input,
613 torch::stable::Tensor const& scale,
614 std::optional<torch::headeronly::IntHeaderOnlyArrayRef> group_shape =
615 std::nullopt);
616
617 void dynamic_scaled_fp8_quant(torch::stable::Tensor& out,
618 torch::stable::Tensor const& input,
619 torch::stable::Tensor& scale);
620
621 void dynamic_per_token_scaled_fp8_quant(
622 torch::stable::Tensor& out, torch::stable::Tensor const& input,
623 torch::stable::Tensor& scale,
624 std::optional<torch::stable::Tensor> const& scale_ub);
625
626 // GPTQ kernels (shared CUDA/ROCm)
627 torch::stable::Tensor gptq_gemm(torch::stable::Tensor a,
628 torch::stable::Tensor b_q_weight,
629 torch::stable::Tensor b_gptq_qzeros,
630 torch::stable::Tensor b_gptq_scales,
631 bool use_exllama, bool use_v2_format,
632 int64_t bit);
633
634 void gptq_shuffle(torch::stable::Tensor q_weight, int64_t bit);
635
636 // Cache ops (shared CUDA/ROCm)
637 void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst,
638 int64_t block_size_in_bytes,
639 const torch::stable::Tensor& block_mapping);
640
641 // Batch swap: submit all block copies in a single driver call.
642 void swap_blocks_batch(const torch::stable::Tensor& src_ptrs,
643 const torch::stable::Tensor& dst_ptrs,
644 const torch::stable::Tensor& sizes,
645 bool is_src_access_order_any);
646
647 void reshape_and_cache(torch::stable::Tensor& key, torch::stable::Tensor& value,
648 torch::stable::Tensor& key_cache,
649 torch::stable::Tensor& value_cache,
650 torch::stable::Tensor& slot_mapping,
651 const std::string& kv_cache_dtype,
652 torch::stable::Tensor& k_scale,
653 torch::stable::Tensor& v_scale);
654
655 void reshape_and_cache_flash(
656 torch::stable::Tensor& key, torch::stable::Tensor& value,
657 torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache,
658 torch::stable::Tensor& slot_mapping, const std::string& kv_cache_dtype,
659 torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale);
660
661 void concat_and_cache_mla(torch::stable::Tensor& kv_c,
662 torch::stable::Tensor& k_pe,
663 torch::stable::Tensor& kv_cache,
664 torch::stable::Tensor& slot_mapping,
665 const std::string& kv_cache_dtype,
666 torch::stable::Tensor& scale);
667
668 void concat_and_cache_mla_grouped(
669 torch::stable::Tensor& kv_c, torch::stable::Tensor& k_pe,
670 torch::stable::Tensor& kv_cache_ptrs, torch::stable::Tensor& slot_mapping,
671 int64_t block_size, int64_t block_stride, int64_t entry_stride,
672 std::optional<torch::stable::Tensor> kv_scales,
673 const std::string& kv_cache_dtype);
674
675 #ifndef USE_ROCM
676 // HiSparse kernels use raw PTX in the row copy; CUDA-only.
677 void hisparse_resolve_residency(
678 torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
679 torch::stable::Tensor const& hot_block_table,
680 torch::stable::Tensor const& global_indices,
681 torch::stable::Tensor& hot_indices,
682 torch::stable::Tensor& device_global_indices,
683 torch::stable::Tensor& lru_slots,
684 std::optional<torch::stable::Tensor> const& request_state_indices,
685 int64_t region_stride, int64_t max_union_rows,
686 std::optional<torch::stable::Tensor> const& miss_mask,
687 std::optional<torch::stable::Tensor> const& stats,
688 std::optional<torch::stable::Tensor> const& attention_indices,
689 int64_t attention_block_stride,
690 std::optional<torch::stable::Tensor> const& request_ids,
691 std::optional<torch::stable::Tensor> const& source_block_table,
692 int64_t source_block_size,
693 std::optional<torch::stable::Tensor> const& resolved_global_indices,
694 std::optional<torch::stable::Tensor> const& valid_counts,
695 std::optional<torch::stable::Tensor> const& swap_host_physical_rows,
696 std::optional<torch::stable::Tensor> const& swap_device_physical_rows,
697 std::optional<torch::stable::Tensor> const& swap_counts,
698 std::optional<torch::stable::Tensor> const& resident_block_table,
699 int64_t resident_block_size, int64_t resident_null_block);
700
701 void hisparse_invalidate_written_slots(
702 torch::stable::Tensor& device_global_indices,
703 torch::stable::Tensor const& request_state_indices,
704 torch::stable::Tensor const& req_id_per_token,
705 torch::stable::Tensor const& written_slots);
706
707 void hisparse_gather_plan(
708 torch::stable::Tensor const& host_cache, torch::stable::Tensor& hot_cache,
709 torch::stable::Tensor const& global_indices,
710 torch::stable::Tensor const& hot_indices,
711 torch::stable::Tensor const& miss_mask,
712 std::optional<torch::stable::Tensor> const& request_state_indices,
713 std::optional<torch::stable::Tensor> const& attention_indices,
714 int64_t attention_block_stride);
715
716 void hisparse_gather_compact(torch::stable::Tensor const& host_cache,
717 torch::stable::Tensor& hot_cache,
718 torch::stable::Tensor const& miss_global_indices,
719 torch::stable::Tensor const& miss_hot_indices,
720 torch::stable::Tensor const& miss_counts);
721
722 #endif // !USE_ROCM
723
724 // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
725 void concat_and_cache_mla_rope_fused(
726 torch::stable::Tensor& positions, torch::stable::Tensor& q_pe,
727 torch::stable::Tensor& k_pe, torch::stable::Tensor& kv_c,
728 torch::stable::Tensor& rope_cos_sin_cache, bool rope_is_neox,
729 torch::stable::Tensor& slot_mapping, torch::stable::Tensor& kv_cache,
730 const std::string& kv_cache_dtype,
731 torch::stable::Tensor& kv_cache_quant_scale);
732
733 // Just for unittest
734 void convert_fp8(torch::stable::Tensor& dst_cache,
735 torch::stable::Tensor& src_cache, const double scale,
736 const std::string& kv_cache_dtype);
737
738 void gather_and_maybe_dequant_cache(
739 torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
740 // ENTRIES...]
741 torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
742 torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
743 torch::stable::Tensor const& cu_seq_lens, // [BATCH+1]
744 torch::stable::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS]
745 int64_t num_tokens, const std::string& kv_cache_dtype,
746 torch::stable::Tensor const& scale,
747 std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
748
749 // TODO(hc): cp_gather_cache need support scaled kvcahe in the future.
750 void cp_gather_cache(
751 torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
752 // ENTRIES...]
753 torch::stable::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
754 torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
755 torch::stable::Tensor const& cu_seq_lens, // [BATCH+1]
756 int64_t batch_size,
757 std::optional<torch::stable::Tensor> seq_starts = std::nullopt);
758
759 // Gather and upconvert FP8 KV cache to BF16 workspace
760 void cp_gather_and_upconvert_fp8_kv_cache(
761 torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
762 // 656]
763 torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
764 torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
765 torch::stable::Tensor const& workspace_starts, // [BATCH]
766 int64_t batch_size,
767 std::optional<torch::stable::Tensor> seq_starts = std::nullopt,
768 std::optional<torch::stable::Tensor> host_cache = std::nullopt,
769 std::optional<torch::stable::Tensor> host_row_ids = std::nullopt,
770 std::optional<torch::stable::Tensor> device_row_ids = std::nullopt);
771
772 // Gather and upconvert an nvfp4_ds_mla KV cache to a BF16 workspace
773 void cp_gather_and_upconvert_nvfp4_kv_cache(
774 torch::stable::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE,
775 // 352]
776 torch::stable::Tensor const& dst, // [TOT_TOKENS, 576]
777 torch::stable::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
778 torch::stable::Tensor const& workspace_starts, // [BATCH]
779 int64_t batch_size);
780
781 // Indexer K quantization and cache function
782 void indexer_k_quant_and_cache(
783 torch::stable::Tensor& k, // [num_tokens, head_dim]
784 torch::stable::Tensor& kv_cache, // [num_blocks, block_size,
785 // cache_stride]
786 torch::stable::Tensor& slot_mapping, // [num_tokens]
787 int64_t quant_block_size, // quantization block size
788 const std::string& scale_fmt);
789
790 // Concatenate query nope and rope for MLA/DSA attention
791 void concat_mla_q(
792 torch::stable::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
793 torch::stable::Tensor& q_pe, // [num_tokens, num_heads, rope_dim]
794 torch::stable::Tensor& q_out); // [num_tokens, num_heads, nope_dim +
795 // rope_dim]
796
797 // Extract function to gather quantized K cache
798 void cp_gather_indexer_k_quant_cache(
799 const torch::stable::Tensor& kv_cache, // [num_blocks, block_size,
800 // cache_stride]
801 torch::stable::Tensor& dst_k, // [num_tokens, head_dim]
802 torch::stable::Tensor& dst_scale, // [num_tokens, head_dim /
803 // quant_block_size * 4]
804 const torch::stable::Tensor& block_table, // [batch_size, num_blocks]
805 const torch::stable::Tensor& cu_seq_lens); // [batch_size + 1]
806
807 // Fused vocab-parallel embedding lookup (see
808 // vocab_parallel_embedding_kernels.cu).
809 void vocab_parallel_embedding(
810 torch::stable::Tensor& out, const torch::stable::Tensor& input_ids,
811 const torch::stable::Tensor& weight, int64_t org_vocab_start_index,
812 int64_t org_vocab_end_index, int64_t num_org_vocab_padding,
813 int64_t added_vocab_start_index, int64_t added_vocab_end_index);
814
815 // LongCat n-gram embedding index kernel (see ngram_embedding_kernels.cu).
816 void ngram_compute_n_gram_ids(
817 int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights,
818 torch::stable::Tensor& ne_mods,
819 torch::stable::Tensor& exclusive_ne_embedder_size_sums,
820 torch::stable::Tensor& exclusive_req_len_sums,
821 torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices,
822 torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids);