# [Experimental] Model Runner V2

This directory contains the new model runner which is under active development.

please see [Model Runner V2](https://github.com/vllm-project/vllm-ascend/issues/5208)
to get specific plans.

## Gaps with vLLM (To Be Addressed)

- [ ] `set_cos_and_sin` & `update_cos_sin`

    Why: 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_cache`

    Why: 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_wrapper`

    Why: 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_workspaces` when each capturing is exactly completed to avoid OOM, thus that leads us to regard `weak_ref_workspaces` as post-processing in `torch.npu.graph` and 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_wrapper`

    Why: 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_wrapper`

    Why: 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`.