vllm-slot

vllm/vllm/v1/worker/block_table.py · ef9339e · L348–376 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

348        start_idx = tl.load(query_start_loc_ptr + req_idx).to(tl.int64)
349        end_idx = tl.load(query_start_loc_ptr + req_idx + 1).to(tl.int64)
350
351        virtual_block_size = KV_CACHE_BLOCK_SIZE * TOTAL_CP_WORLD_SIZE
352        row_offset = req_idx * block_table_stride
353        for i in range(start_idx, end_idx, BLOCK_SIZE):
354            offsets = i + tl.arange(0, BLOCK_SIZE)
355            mask = offsets < end_idx
356            pos = tl.load(positions_ptr + offsets, mask=mask, other=0)
357            virtual_block_indices = pos // virtual_block_size
358            virtual_block_offsets = pos - virtual_block_indices * virtual_block_size
359            is_local = (
360                virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE
361            ) % TOTAL_CP_WORLD_SIZE == TOTAL_CP_RANK
362            local_block_offsets = (
363                virtual_block_offsets
364                // (TOTAL_CP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE)
365            ) * CP_KV_CACHE_INTERLEAVE_SIZE + (
366                virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE
367            )
368
369            block_numbers = tl.load(
370                block_table_ptr + row_offset + virtual_block_indices,
371                mask=mask & is_local,
372                other=0,
373            ).to(tl.int64)
374            slot_ids = block_numbers * KV_CACHE_BLOCK_SIZE + local_block_offsets
375            slot_ids = tl.where(is_local, slot_ids, PAD_ID)
376            tl.store(slot_mapping_ptr + offsets, slot_ids, mask=mask)