Separate allocation logic from scheduler

[cctry]

Motivation

Preparation for mem_cache V2.
The allocation logic is moved to mem_cache/ from schedule_batch.py. Ideally, scheduling code should only interact with tree_cache in V1 and memory_manager in V2.

Modifications

• create mem_cache/common.py for allocation functions operating allocator and tree cache
• refactor the allocation to use two wrappers alloc_for_extend and alloc_for_decode
  • in prepare_for_decode, the increment of seqlen is moved after alloc_for_decode for clarity
  • in prepare_for_extend, some allocation-needed fields are set before alloc_for_extend
• for spec decode, the allocation functions are unchanged and they are to be changed after spec V2
• in bench_one_batch.py, create a dummy tree_cache as the placeholder

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.

[xiezhq-hermann]

why evicting proactively here?

[cctry]

this is actually evict_if_needed

         if self.token_to_kv_pool_allocator.available_size() < num_tokens:
                if self.tree_cache is not None:
                    self.tree_cache.evict(num_tokens)

[xiezhq-hermann]

maybe we can rename it a bit, it's actually check for availability and evict if needed

[cctry]

what will be the case that we want to evict nodes regardless of the availability

i also feel the eviction policy is something non-trivial so evict_from_tree_cache will encapsulate complex logic from upper code

[xiezhq-hermann]

the current self.tree_cache.evict is indeed to evict nodes to meet the requirement regardless of the availability, and I think we can probably keep eviction policy under the hood too like what we have for now

[airMeng]

Hi, I find batch,token_to_kv_pool_allocator.page_size not always equals to batch.tree_cache.page_size, then the paged config will go to the wrong allocation, breaks a lot of cases

[cctry]

could you let me know when the page sizes are different? we can add tests to capture this in the future

[airMeng]

#11313 (comment) seems the page sizes should be the same and this is just a bug.

But the question is why we need to access the same page size from different places? echos the suggestions here #11645 (comment)

[airMeng]

why hard-code to 1 instead page_size=model_runner.server_args.page_size ?

[cctry]

Thanks for pointing out. wIll fix this