[Performance]: Custom AscendC Kernel of Multi-Step Prepare Input (#814)
### What this PR does / why we need it? - According to https://github.com/vllm-project/vllm-ascend/issues/807, we pull request for customer ascendc kernel of multi-step. - also a bug we found in multi_step_runner.py is fixed when we use multi-step on V0 Engine. ### Does this PR introduce _any_ user-facing change? no user-facing change ### How was this patch tested? we add Unit Test file and offline inference file to test the custom ascendc kernel. See test/ops/test_multi_step.py and examples/offline_multi_step.py --------- Signed-off-by: wan_danfeng <wonderful199082@126.com>
This commit is contained in:
12
csrc/ops.h
12
csrc/ops.h
@@ -46,4 +46,16 @@ namespace vllm_ascend {
|
||||
auto new_tensor = at_npu::native::from_blob(data_ptr, sizes, strides, options);
|
||||
return new_tensor;
|
||||
}
|
||||
extern void launch_advance_step_flashattn(
|
||||
void* stream,
|
||||
int64_t num_seqs,
|
||||
int64_t num_queries,
|
||||
int64_t block_size,
|
||||
int64_t* input_tokens_ptr,
|
||||
int64_t* sampled_token_ids_ptr,
|
||||
int64_t* input_positions_ptr,
|
||||
int32_t* seq_lens_ptr,
|
||||
int32_t* slot_mapping_ptr,
|
||||
int32_t* block_tables_ptr,
|
||||
int64_t block_tables_stride);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user