Skip to content

Commit a90b7cd

Browse files
Bruce-Lee-LYBruce-Lee-LY
committed
fix: ci test
1 parent 337c97a commit a90b7cd

File tree

10 files changed

+425
-240
lines changed

10 files changed

+425
-240
lines changed

‎rtp_llm/cpp/devices/DeviceBase.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -167,6 +167,9 @@ class DeviceBase: public DeviceOps {
167167
torch::Tensor ropeCache() const {
168168
return rope_cache_;
169169
}
170+
int ropeCacheDim() const {
171+
return rope_cache_dim_;
172+
}
170173

171174
public:
172175
// device-independence op implementations
@@ -228,6 +231,7 @@ class DeviceBase: public DeviceOps {
228231
std::once_flag rope_cache_flag_;
229232
bool use_rope_cache_ = false;
230233
torch::Tensor rope_cache_;
234+
int rope_cache_dim_;
231235

232236
protected:
233237
std::unique_ptr<BufferManager> buffer_manager_;

‎rtp_llm/cpp/devices/base_tests/AttentionOpTest.hpp‎

Lines changed: 123 additions & 84 deletions
Large diffs are not rendered by default.

‎rtp_llm/cpp/devices/base_tests/UnfusedAttentionTest.hpp‎

Lines changed: 40 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -291,7 +291,8 @@ void UnfusedAttentionTest::addFusedQKVBiasTransposeTest(size_t batch_size,
291291
params.common.cu_seqlens->data<int>(),
292292
params.common.cu_seqlens_without_prefix->data<int>(),
293293
device->use_rope_cache_,
294-
device->use_rope_cache_ && device->rope_cache_.defined() ?
294+
device->use_rope_cache_ && device->rope_cache_.defined()
295+
&& device->rope_cache_dim_ == params.configs.rope_config.dim ?
295296
device->rope_cache_.data_ptr<float>() :
296297
nullptr,
297298
batch_size,
@@ -335,7 +336,8 @@ void UnfusedAttentionTest::addFusedQKVBiasTransposeTest(size_t batch_size,
335336
params.common.cu_seqlens->data<int>(),
336337
params.common.cu_seqlens_without_prefix->data<int>(),
337338
device->use_rope_cache_,
338-
device->use_rope_cache_ && device->rope_cache_.defined() ?
339+
device->use_rope_cache_ && device->rope_cache_.defined()
340+
&& device->rope_cache_dim_ == params.configs.rope_config.dim ?
339341
device->rope_cache_.data_ptr<float>() :
340342
nullptr,
341343
batch_size,
@@ -374,40 +376,42 @@ void UnfusedAttentionTest::addFusedQKVBiasTransposeTest(size_t batch_size,
374376
bool store_kv = true;
375377
bool store_cache = false;
376378

377-
DISPATCH_CUDA_FUNCTION_DATA_TYPE(
378-
params.input.type(),
379-
invokeAddFusedQKVBiasTranspose,
380-
q_no_transpose_output->data(),
381-
q_output->data(),
382-
k_output->data(),
383-
v_output->data(),
384-
&prefix_prompt_param,
385-
params.input.data(),
386-
qkv_buf_fp8 ? qkv_buf_fp8->data() : nullptr,
387-
params.common.position_ids->data<int>(),
388-
params.weights.qkv_weight->bias->data(),
389-
params.common.padding_offset->data<int>(),
390-
params.common.cu_seqlens->data<int>(),
391-
params.common.cu_seqlens_without_prefix->data<int>(),
392-
device->use_rope_cache_,
393-
device->use_rope_cache_ && device->rope_cache_.defined() ? device->rope_cache_.data_ptr<float>() : nullptr,
394-
batch_size,
395-
seq_len,
396-
token_num,
397-
num_heads,
398-
num_key_value_heads,
399-
head_dim,
400-
params.configs.rope_config,
401-
params.configs.use_logn_attn,
402-
scale_out_ptr,
403-
int8_mode,
404-
use_paged_fmha,
405-
store_qkv,
406-
store_q_no_transpose,
407-
store_q,
408-
store_kv,
409-
store_cache,
410-
device->getStream());
379+
DISPATCH_CUDA_FUNCTION_DATA_TYPE(params.input.type(),
380+
invokeAddFusedQKVBiasTranspose,
381+
q_no_transpose_output->data(),
382+
q_output->data(),
383+
k_output->data(),
384+
v_output->data(),
385+
&prefix_prompt_param,
386+
params.input.data(),
387+
qkv_buf_fp8 ? qkv_buf_fp8->data() : nullptr,
388+
params.common.position_ids->data<int>(),
389+
params.weights.qkv_weight->bias->data(),
390+
params.common.padding_offset->data<int>(),
391+
params.common.cu_seqlens->data<int>(),
392+
params.common.cu_seqlens_without_prefix->data<int>(),
393+
device->use_rope_cache_,
394+
device->use_rope_cache_ && device->rope_cache_.defined()
395+
&& device->rope_cache_dim_ == params.configs.rope_config.dim ?
396+
device->rope_cache_.data_ptr<float>() :
397+
nullptr,
398+
batch_size,
399+
seq_len,
400+
token_num,
401+
num_heads,
402+
num_key_value_heads,
403+
head_dim,
404+
params.configs.rope_config,
405+
params.configs.use_logn_attn,
406+
scale_out_ptr,
407+
int8_mode,
408+
use_paged_fmha,
409+
store_qkv,
410+
store_q_no_transpose,
411+
store_q,
412+
store_kv,
413+
store_cache,
414+
device->getStream());
411415

412416
device->syncAndCheck();
413417

‎rtp_llm/cpp/devices/cuda_impl/CudaAttentionOp.cc‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,8 @@ void CudaDevice::getRopeCacheOnce(const RopeConfig& rope_config, int max_positio
9999
std::call_once(rope_cache_flag_, [&]() {
100100
use_rope_cache_ = rope_config.style == RopeStyle::Base || rope_config.style == RopeStyle::Yarn;
101101
if (use_rope_cache_) {
102-
rope_cache_ = getRopeCache(rope_config, max_position_embeddings);
102+
rope_cache_ = getRopeCache(rope_config, max_position_embeddings);
103+
rope_cache_dim_ = rope_config.dim;
103104
}
104105
});
105106
}
@@ -219,7 +220,9 @@ AttentionModuleOutput CudaDevice::contextAttention(const AttentionModuleParams&
219220
params.common.cu_seqlens->data<int>(),
220221
params.common.cu_seqlens_without_prefix->data<int>(),
221222
use_rope_cache_,
222-
use_rope_cache_ && rope_cache_.defined() ? rope_cache_.data_ptr<float>() : nullptr,
223+
use_rope_cache_ && rope_cache_.defined() && rope_cache_dim_ == params.configs.rope_config.dim ?
224+
rope_cache_.data_ptr<float>() :
225+
nullptr,
223226
batch_size,
224227
seq_len,
225228
token_num,
@@ -410,8 +413,10 @@ AttentionModuleOutput CudaDevice::decoderSelfAttention(const AttentionModulePara
410413
params.weights.qkv_weight->bias->data() :
411414
nullptr,
412415
use_rope_cache_,
413-
use_rope_cache_ && rope_cache_.defined() ? rope_cache_.data_ptr<float>() :
414-
nullptr,
416+
use_rope_cache_ && rope_cache_.defined()
417+
&& rope_cache_dim_ == params.configs.rope_config.dim ?
418+
rope_cache_.data_ptr<float>() :
419+
nullptr,
415420
batch_size,
416421
local_head_num,
417422
local_kv_head_num,

‎rtp_llm/cpp/devices/rocm_impl/ROCmAttentionOp.cc‎

Lines changed: 41 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -571,7 +571,8 @@ void ROCmDevice::getRopeCacheOnce(const RopeConfig& rope_config, int max_positio
571571
std::call_once(rope_cache_flag_, [&]() {
572572
use_rope_cache_ = rope_config.style == RopeStyle::Base;
573573
if (use_rope_cache_) {
574-
rope_cache_ = getRopeCache(rope_config, max_position_embeddings);
574+
rope_cache_ = getRopeCache(rope_config, max_position_embeddings);
575+
rope_cache_dim_ = rope_config.dim;
575576
}
576577
});
577578
}
@@ -778,42 +779,46 @@ AttentionModuleOutput ROCmDevice::contextAttention(const AttentionModuleParams&
778779
}
779780
check_cuda_error();
780781
} else {
781-
DISPATCH_CUDA_FUNCTION_DATA_TYPE(datatype,
782-
invokeAddFusedQKVBiasTranspose,
783-
nullptr,
784-
q_output->data(),
785-
k_output->data(),
786-
v_output->data(),
787-
&prefix_prompt_param,
788-
params.input.data(),
782+
DISPATCH_CUDA_FUNCTION_DATA_TYPE(
783+
datatype,
784+
invokeAddFusedQKVBiasTranspose,
785+
nullptr,
786+
q_output->data(),
787+
k_output->data(),
788+
v_output->data(),
789+
&prefix_prompt_param,
790+
params.input.data(),
791+
nullptr,
792+
params.common.position_ids ? params.common.position_ids->dataWithOffset<int>(
793+
decoder_batch_size * params.configs.rope_config.index_factor) :
789794
nullptr,
790-
params.common.position_ids ?
791-
params.common.position_ids->dataWithOffset<int>(
792-
decoder_batch_size * params.configs.rope_config.index_factor) :
793-
nullptr,
794-
params.configs.fuse_qkv_add_bias && params.weights.qkv_weight->bias ?
795-
params.weights.qkv_weight->bias->data() :
796-
nullptr,
797-
params.common.padding_offset->data<int>(),
798-
params.common.cu_seqlens->data<int>(),
799-
params.common.cu_seqlens_without_prefix->data<int>(),
800-
batch_size,
801-
seq_len,
802-
token_num,
803-
head_num,
804-
kv_head_num,
805-
size_per_head,
806-
params.configs.rope_config,
807-
params.configs.use_logn_attn,
808-
scale_out_ptr,
809-
int8_mode,
810-
false,
811-
store_qkv,
812-
false,
813-
store_q,
814-
store_kv,
815-
store_cache,
816-
stream_);
795+
params.configs.fuse_qkv_add_bias && params.weights.qkv_weight->bias ?
796+
params.weights.qkv_weight->bias->data() :
797+
nullptr,
798+
params.common.padding_offset->data<int>(),
799+
params.common.cu_seqlens->data<int>(),
800+
params.common.cu_seqlens_without_prefix->data<int>(),
801+
use_rope_cache_,
802+
use_rope_cache_ && rope_cache_.defined() && rope_cache_dim_ == params.configs.rope_config.dim ?
803+
rope_cache_.data_ptr<float>() :
804+
nullptr,
805+
batch_size,
806+
seq_len,
807+
token_num,
808+
head_num,
809+
kv_head_num,
810+
size_per_head,
811+
params.configs.rope_config,
812+
params.configs.use_logn_attn,
813+
scale_out_ptr,
814+
int8_mode,
815+
false,
816+
store_qkv,
817+
false,
818+
store_q,
819+
store_kv,
820+
store_cache,
821+
stream_);
817822
check_cuda_error();
818823
}
819824
writeCacheStore(params);

0 commit comments

Comments
 (0)