[Experimental] Model Runner V2
This directory contains the new model runner which is under active development.
please see Model Runner V2 to get specific plans.
Gaps with vLLM (To Be Addressed)
-
set_cos_and_sin&update_cos_sinWhy: DeepSeek-like models (mla) still need cos/sin setting and updating in model_runner. These should be removed when mla can solve cos/sin internally.
Location:
NPUModelRunner.__init__,NPUModelRunner.prepare_inputs,AscendInputBatch.make_dummy. -
_allocate_kv_cache&_reshape_kv_cacheWhy: KV cache requires continuous space (thus divided as K cache and V cache separately) and PD disaggregation requires 2M-aligned tensors for KV cache, so custom KV cache initialization is needed. These should be removed when the above 2 requirements are no longer needed.
Location:
attn_utils._get_layer_kv_cache_specs,attn_utils._get_attention_kv_cache_dims,attn_utils._align_memory,attn_utils._allocate_kv_cache,attn_utils._reshape_kv_cache. -
torch_npu_graph_wrapperWhy: FIA ops in FULL mode need explicit workspace allocating, and each workspace corresponding to each graph (a specific batch_size) should be released via
weak_ref_workspaceswhen each capturing is exactly completed to avoid OOM, thus that leads us to regardweak_ref_workspacesas post-processing intorch.npu.graphand patch it. This should be removed when we don't need such special operations.Location:
utils.torch_cuda_wrapper,utils.torch_npu_graph_wrapper. -
model_runner.graph_manager_wrapperWhy: ModelAclGraphManager needs model_runner's input_buffers and model_state.attn_metadata to update_full_graph_params, so model_runner should be passed into init of ModelAclGraphManager.
Location:
model_runner.NPUModelRunner.initialize_kv_cache. -
speculator.graph_manager_wrapperWhy: EagleAclGraphManager needs speculator's input_buffers and model_state.attn_metadata to update_full_graph_params, so speculator should be passed into __init __ of EagleAclGraphManager.
Location:
speculator.AscendEagleSpeculator.init_cudagraph_manager.