@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 93 lines 4.4 KB
Raw
1 #pragma once
2
3 #include <torch/all.h>
4 #include <c10/util/Optional.h>
5
6 #include <map>
7 #include <vector>
8
9 void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
10 int64_t block_size_in_bytes,
11 const torch::Tensor& block_mapping);
12
13 void swap_blocks_batch(const torch::Tensor& src_ptrs,
14 const torch::Tensor& dst_ptrs,
15 const torch::Tensor& sizes,
16 bool is_src_access_order_any);
17
18 void reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
19 torch::Tensor& key_cache, torch::Tensor& value_cache,
20 torch::Tensor& slot_mapping,
21 const std::string& kv_cache_dtype,
22 torch::Tensor& k_scale, torch::Tensor& v_scale);
23
24 void reshape_and_cache_flash(torch::Tensor& key, torch::Tensor& value,
25 torch::Tensor& key_cache,
26 torch::Tensor& value_cache,
27 torch::Tensor& slot_mapping,
28 const std::string& kv_cache_dtype,
29 torch::Tensor& k_scale, torch::Tensor& v_scale);
30
31 void concat_and_cache_mla(torch::Tensor& kv_c, torch::Tensor& k_pe,
32 torch::Tensor& kv_cache, torch::Tensor& slot_mapping,
33 const std::string& kv_cache_dtype,
34 torch::Tensor& scale);
35
36 // NOTE: k_pe and kv_c order is flipped compared to concat_and_cache_mla
37 void concat_and_cache_mla_rope_fused(
38 torch::Tensor& positions, torch::Tensor& q_pe, torch::Tensor& k_pe,
39 torch::Tensor& kv_c, torch::Tensor& rope_cos_sin_cache, bool rope_is_neox,
40 torch::Tensor& kv_cache_slot_mapping, torch::Tensor& kv_cache,
41 const std::string& kv_cache_dtype, torch::Tensor& kv_cache_quant_scale);
42
43 // Just for unittest
44 void convert_fp8(torch::Tensor& dst_cache, torch::Tensor& src_cache,
45 const double scale, const std::string& kv_cache_dtype);
46
47 void gather_and_maybe_dequant_cache(
48 torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
49 torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
50 torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
51 torch::Tensor const& cu_seq_lens, // [BATCH+1]
52 torch::Tensor const& token_to_seq, // [MAX_TOKEN_ACROSS_CHUNKS]
53 int64_t num_tokens, const std::string& kv_cache_dtype,
54 torch::Tensor const& scale,
55 std::optional<torch::Tensor> seq_starts = std::nullopt);
56
57 // TODO(hc): cp_gather_cache need support scaled kvcahe in the future.
58 void cp_gather_cache(
59 torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
60 torch::Tensor const& dst, // [TOT_TOKENS, ENTRIES...]
61 torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
62 torch::Tensor const& cu_seq_lens, // [BATCH+1]
63 int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
64
65 // Gather and upconvert FP8 KV cache to BF16 workspace
66 void cp_gather_and_upconvert_fp8_kv_cache(
67 torch::Tensor const& src_cache, // [NUM_BLOCKS, BLOCK_SIZE, 656]
68 torch::Tensor const& dst, // [TOT_TOKENS, 576]
69 torch::Tensor const& block_table, // [BATCH, BLOCK_INDICES]
70 torch::Tensor const& workspace_starts, // [BATCH]
71 int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
72
73 // Indexer K quantization and cache function
74 void indexer_k_quant_and_cache(
75 torch::Tensor& k, // [num_tokens, head_dim]
76 torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
77 torch::Tensor& slot_mapping, // [num_tokens]
78 int64_t quant_block_size, // quantization block size
79 const std::string& scale_fmt);
80
81 // Concatenate query nope and rope for MLA/DSA attention
82 void concat_mla_q(
83 torch::Tensor& ql_nope, // [num_tokens, num_heads, nope_dim]
84 torch::Tensor& q_pe, // [num_tokens, num_heads, rope_dim]
85 torch::Tensor& q_out); // [num_tokens, num_heads, nope_dim + rope_dim]
86
87 // Extract function to gather quantized K cache
88 void cp_gather_indexer_k_quant_cache(
89 const torch::Tensor& kv_cache, // [num_blocks, block_size, cache_stride]
90 torch::Tensor& dst_k, // [num_tokens, head_dim]
91 torch::Tensor& dst_scale, // [num_tokens, head_dim / quant_block_size * 4]
92 const torch::Tensor& block_table, // [batch_size, num_blocks]
93 const torch::Tensor& cu_seq_lens); // [batch_size + 1]