Skip to content

Commit 6b9cfa6

Browse files
whjthuwhjthu
committed
refactor: simplify paged graph eligibility
1 parent 768ebea commit 6b9cfa6

1 file changed

Lines changed: 63 additions & 133 deletions

File tree

‎csrc/engine/compiler/paged_compiler.cpp‎

Lines changed: 63 additions & 133 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
#include "../../utils.hpp"
88

99
#include <algorithm>
10-
#include <array>
1110
#include <cstdint>
1211
#include <cstdlib>
1312
#include <initializer_list>
@@ -34,123 +33,6 @@ constexpr size_t kBaichuanFixedPrefillBatchSize = 1;
3433
constexpr size_t kBaichuanFixedPrefillSequenceLength = 10;
3534
constexpr size_t kBaichuanFixedPrefillBlockSize = 256;
3635

37-
struct ReviewedPagedGraphProfile {
38-
std::string_view model_type;
39-
size_t hidden_size;
40-
size_t num_attention_heads;
41-
size_t num_key_value_heads;
42-
size_t head_dim;
43-
size_t block_size;
44-
size_t minimum_num_blocks;
45-
std::optional<size_t> exact_num_blocks;
46-
std::optional<size_t> max_batch_size;
47-
std::optional<size_t> tensor_parallel_world_size;
48-
std::optional<size_t> num_hidden_layers;
49-
std::optional<size_t> position_id_axes;
50-
std::optional<infinilm::runtime::DataType> dtype;
51-
bool require_unquantized;
52-
};
53-
54-
struct PagedGraphProperties {
55-
infinilm::runtime::Device::Type device_type;
56-
backends::AttentionBackend attention_backend;
57-
size_t tensor_parallel_world_size;
58-
size_t block_size;
59-
size_t num_blocks;
60-
size_t max_batch_size;
61-
std::string model_type;
62-
size_t hidden_size;
63-
size_t num_hidden_layers;
64-
size_t num_attention_heads;
65-
size_t num_key_value_heads;
66-
size_t head_dim;
67-
size_t position_id_axes;
68-
infinilm::runtime::DataType dtype;
69-
quantization::QuantScheme quant_scheme;
70-
quantization::KVQuantAlgo kv_quant_scheme;
71-
};
72-
73-
const ReviewedPagedGraphProfile kBaichuanFixedPrefillProfile{
74-
"baichuan", 4096, 32, 32, 128, kBaichuanFixedPrefillBlockSize, 1,
75-
1, 1, 2, 32, 1, std::nullopt, true};
76-
77-
const std::array<ReviewedPagedGraphProfile, 2> kShortDecodeProfiles{{
78-
{"internlm3", 4096, 32, 2, 128, kShortDecodeBlockSize,
79-
kShortDecodeBlockTableWidth, std::nullopt, std::nullopt,
80-
std::nullopt, std::nullopt, std::nullopt, std::nullopt, false},
81-
{"chatglm", 4096, 32, 2, 128, kShortDecodeBlockSize,
82-
kShortDecodeBlockTableWidth, 512, std::nullopt, 1, 28, 1,
83-
infinilm::runtime::DataType::kFloat16, true},
84-
}};
85-
86-
template <typename T>
87-
bool matches_optional_constraint(
88-
const T &actual,
89-
const std::optional<T> &expected) {
90-
return !expected.has_value() || actual == expected.value();
91-
}
92-
93-
std::optional<PagedGraphProperties> read_paged_graph_properties(
94-
const cache::PagedKVCacheConfig &paged_config,
95-
const config::ModelConfig *model_config,
96-
infinilm::runtime::Device::Type device_type,
97-
const PagedGraphConfig &graph_config) {
98-
graph_config.validate();
99-
if (model_config == nullptr) {
100-
return std::nullopt;
101-
}
102-
103-
const size_t hidden_size = model_config->get_or<size_t>("hidden_size", 0);
104-
const size_t num_attention_heads = model_config->get_or<size_t>("num_attention_heads", 0);
105-
const size_t head_dim = model_config->get_or<size_t>(
106-
"head_dim",
107-
num_attention_heads == 0 ? 0 : hidden_size / num_attention_heads);
108-
109-
return PagedGraphProperties{
110-
device_type,
111-
graph_config.attention_backend,
112-
static_cast<size_t>(graph_config.tp_size),
113-
paged_config.block_size(),
114-
paged_config.num_blocks(),
115-
paged_config.max_batch_size(),
116-
model_config->get_or<std::string>("model_type", ""),
117-
hidden_size,
118-
model_config->get_or<size_t>("num_hidden_layers", 0),
119-
num_attention_heads,
120-
model_config->get_or<size_t>("num_key_value_heads", 0),
121-
head_dim,
122-
model_config->get_or<size_t>("position_id_axes", 1),
123-
model_config->get_dtype(),
124-
model_config->get_quant_scheme(),
125-
model_config->get_kv_quant_scheme(),
126-
};
127-
}
128-
129-
bool matches_reviewed_paged_graph_profile(
130-
const PagedGraphProperties &properties,
131-
const ReviewedPagedGraphProfile &profile) {
132-
return properties.device_type == infinilm::runtime::Device::Type::kNvidia
133-
&& properties.attention_backend == backends::AttentionBackend::FLASH_ATTN
134-
&& properties.model_type == profile.model_type
135-
&& properties.hidden_size == profile.hidden_size
136-
&& properties.num_attention_heads == profile.num_attention_heads
137-
&& properties.num_key_value_heads == profile.num_key_value_heads
138-
&& properties.head_dim == profile.head_dim
139-
&& properties.block_size == profile.block_size
140-
&& properties.num_blocks >= profile.minimum_num_blocks
141-
&& matches_optional_constraint(properties.num_blocks, profile.exact_num_blocks)
142-
&& matches_optional_constraint(properties.max_batch_size, profile.max_batch_size)
143-
&& matches_optional_constraint(
144-
properties.tensor_parallel_world_size,
145-
profile.tensor_parallel_world_size)
146-
&& matches_optional_constraint(properties.num_hidden_layers, profile.num_hidden_layers)
147-
&& matches_optional_constraint(properties.position_id_axes, profile.position_id_axes)
148-
&& matches_optional_constraint(properties.dtype, profile.dtype)
149-
&& (!profile.require_unquantized
150-
|| (properties.quant_scheme == quantization::QuantScheme::NONE
151-
&& properties.kv_quant_scheme == quantization::KVQuantAlgo::NONE));
152-
}
153-
15436
bool has_mamba_cache(const infinilm::global_state::ForwardContext &forward_context) {
15537
auto has_state = [](const std::vector<infinilm::runtime::Tensor> &state_vec) {
15638
for (const auto &state : state_vec) {
@@ -179,11 +61,35 @@ bool supports_baichuan_fixed_prefill_graph(
17961
return false;
18062
}
18163

182-
const auto properties = read_paged_graph_properties(
183-
paged_config, model_config, device_type, graph_config);
184-
return properties.has_value()
185-
&& matches_reviewed_paged_graph_profile(
186-
properties.value(), kBaichuanFixedPrefillProfile);
64+
graph_config.validate();
65+
if (model_config == nullptr) {
66+
return false;
67+
}
68+
69+
const size_t hidden_size = model_config->get_or<size_t>("hidden_size", 0);
70+
const size_t num_attention_heads = model_config->get_or<size_t>("num_attention_heads", 0);
71+
const size_t head_dim = model_config->get_or<size_t>(
72+
"head_dim",
73+
num_attention_heads == 0 ? 0 : hidden_size / num_attention_heads);
74+
(void)model_config->get_dtype();
75+
const auto quant_scheme = model_config->get_quant_scheme();
76+
const auto kv_quant_scheme = model_config->get_kv_quant_scheme();
77+
78+
return device_type == infinilm::runtime::Device::Type::kNvidia
79+
&& graph_config.attention_backend == backends::AttentionBackend::FLASH_ATTN
80+
&& model_config->get_or<std::string>("model_type", "") == "baichuan"
81+
&& hidden_size == 4096
82+
&& num_attention_heads == 32
83+
&& model_config->get_or<size_t>("num_key_value_heads", 0) == 32
84+
&& head_dim == 128
85+
&& paged_config.block_size() == kBaichuanFixedPrefillBlockSize
86+
&& paged_config.num_blocks() == 1
87+
&& paged_config.max_batch_size() == 1
88+
&& graph_config.tp_size == 2
89+
&& model_config->get_or<size_t>("num_hidden_layers", 0) == 32
90+
&& model_config->get_or<size_t>("position_id_axes", 1) == 1
91+
&& quant_scheme == quantization::QuantScheme::NONE
92+
&& kv_quant_scheme == quantization::KVQuantAlgo::NONE;
18793
}
18894

18995
bool is_cpu_contiguous_tensor(
@@ -279,16 +185,40 @@ bool supports_reviewed_short_decode_graph(
279185
const config::ModelConfig *model_config,
280186
infinilm::runtime::Device::Type device_type,
281187
const PagedGraphConfig &graph_config) {
282-
const auto properties = read_paged_graph_properties(
283-
paged_config, model_config, device_type, graph_config);
284-
return properties.has_value()
285-
&& std::any_of(
286-
kShortDecodeProfiles.begin(),
287-
kShortDecodeProfiles.end(),
288-
[&](const ReviewedPagedGraphProfile &profile) {
289-
return matches_reviewed_paged_graph_profile(
290-
properties.value(), profile);
291-
});
188+
graph_config.validate();
189+
if (model_config == nullptr) {
190+
return false;
191+
}
192+
193+
const size_t hidden_size = model_config->get_or<size_t>("hidden_size", 0);
194+
const size_t num_attention_heads = model_config->get_or<size_t>("num_attention_heads", 0);
195+
const size_t head_dim = model_config->get_or<size_t>(
196+
"head_dim",
197+
num_attention_heads == 0 ? 0 : hidden_size / num_attention_heads);
198+
const auto dtype = model_config->get_dtype();
199+
const auto quant_scheme = model_config->get_quant_scheme();
200+
const auto kv_quant_scheme = model_config->get_kv_quant_scheme();
201+
202+
if (device_type != infinilm::runtime::Device::Type::kNvidia
203+
|| graph_config.attention_backend != backends::AttentionBackend::FLASH_ATTN
204+
|| hidden_size != 4096
205+
|| num_attention_heads != 32
206+
|| model_config->get_or<size_t>("num_key_value_heads", 0) != 2
207+
|| head_dim != 128
208+
|| paged_config.block_size() != kShortDecodeBlockSize
209+
|| paged_config.num_blocks() < kShortDecodeBlockTableWidth) {
210+
return false;
211+
}
212+
const auto model_type = model_config->get_or<std::string>("model_type", "");
213+
return model_type == "internlm3"
214+
|| (model_type == "chatglm"
215+
&& paged_config.num_blocks() == 512
216+
&& graph_config.tp_size == 1
217+
&& model_config->get_or<size_t>("num_hidden_layers", 0) == 28
218+
&& model_config->get_or<size_t>("position_id_axes", 1) == 1
219+
&& dtype == infinilm::runtime::DataType::kFloat16
220+
&& quant_scheme == quantization::QuantScheme::NONE
221+
&& kv_quant_scheme == quantization::KVQuantAlgo::NONE);
292222
}
293223

294224
bool tensors_compatible(const infinilm::runtime::Tensor &target,

0 commit comments

Comments
 (0)