xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-09 23:07:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

misc merge fixes

Browse Source 
Signed-off-by: Sage Moore <sage@neuralmagic.com>
This commit is contained in:
Sage Moore 2025-07-25 19:47:33 +00:00
parent ee70ce0e4e
commit b9ad5e4588
4 changed files with 24 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
vllm/v1/attention/backends/mla/common.py
View File

@ -563,8 +563,10 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
scheduler_output, scheduler_output,
decode_threshold=1) decode_threshold=1)
def _build_decode(self, block_table_tensor: torch.Tensor, def _build_decode(self,
seq_lens: torch.Tensor): block_table_tensor: torch.Tensor,
seq_lens: torch.Tensor,
ubatch_id: Optional[int] = None):
return MLACommonDecodeMetadata( return MLACommonDecodeMetadata(
block_table=block_table_tensor, block_table=block_table_tensor,
seq_lens=seq_lens, seq_lens=seq_lens,
@ -597,7 +599,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
def build(self, def build(self,
common_prefix_len: int, common_prefix_len: int,
common_attn_metadata: CommonAttentionMetadata, common_attn_metadata: CommonAttentionMetadata,
fast_build: bool = False) -> M: fast_build: bool = False,
ubatch_id: Optional[int] = None) -> M:
num_reqs = common_attn_metadata.num_reqs num_reqs = common_attn_metadata.num_reqs
num_tokens = common_attn_metadata.num_actual_tokens num_tokens = common_attn_metadata.num_actual_tokens
max_query_len = common_attn_metadata.max_query_len max_query_len = common_attn_metadata.max_query_len
@ -720,7 +723,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
decode_metadata = self._build_decode( decode_metadata = self._build_decode(
block_table_tensor=block_table_tensor[:num_decodes, ...], block_table_tensor=block_table_tensor[:num_decodes, ...],
seq_lens=seq_lens[:num_decodes], seq_lens=seq_lens[:num_decodes],
) ubatch_id=ubatch_id)
attn_metadata = self.metadata_cls( attn_metadata = self.metadata_cls(
num_reqs=common_attn_metadata.num_reqs, num_reqs=common_attn_metadata.num_reqs,

7
vllm/v1/attention/backends/mla/flashmla.py
View File

@ -67,8 +67,11 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
self.cg_buf_tile_scheduler_metadata = None self.cg_buf_tile_scheduler_metadata = None
self.cg_buf_num_splits = None self.cg_buf_num_splits = None
def _build_decode(self, block_table_tensor: torch.Tensor, def _build_decode(
seq_lens: torch.Tensor) -> FlashMLADecodeMetadata: self,
block_table_tensor: torch.Tensor,
seq_lens: torch.Tensor,
ubatch_id: Optional[int] = None) -> FlashMLADecodeMetadata:
tile_scheduler_metadata, num_splits = \ tile_scheduler_metadata, num_splits = \
get_mla_metadata( get_mla_metadata(
seq_lens, seq_lens,

7
vllm/v1/attention/backends/utils.py
View File

@ -126,7 +126,7 @@ def _make_metadata_with_slice(
def split_attn_metadata( def split_attn_metadata(
ubatch_slices: list[UbatchSlice], ubatch_slices: list[tuple[slice, slice]],
common_attn_metadata: CommonAttentionMetadata, common_attn_metadata: CommonAttentionMetadata,
) -> list[CommonAttentionMetadata]: ) -> list[CommonAttentionMetadata]:
""" """
@ -136,8 +136,9 @@ def split_attn_metadata(
""" """
results = [] results = []
for ubatch_slice in ubatch_slices: for ubatch_slice in ubatch_slices:
results.append( s = UbatchSlice(request_slice=ubatch_slice[0],
_make_metadata_with_slice(ubatch_slice, common_attn_metadata)) token_slice=ubatch_slice[1])
results.append(_make_metadata_with_slice(s, common_attn_metadata))
return results return results

17
vllm/v1/worker/gpu_model_runner.py
View File

@ -52,7 +52,7 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
from vllm.v1.attention.backends.utils import ( from vllm.v1.attention.backends.utils import (
AttentionMetadataBuilder, CommonAttentionMetadata, AttentionMetadataBuilder, CommonAttentionMetadata,
make_local_attention_virtual_batches) make_local_attention_virtual_batches, split_attn_metadata)
from vllm.v1.core.encoder_cache_manager import compute_encoder_budget from vllm.v1.core.encoder_cache_manager import compute_encoder_budget
from vllm.v1.kv_cache_interface import (AttentionSpec, from vllm.v1.kv_cache_interface import (AttentionSpec,
ChunkedLocalAttentionSpec, ChunkedLocalAttentionSpec,
@ -878,17 +878,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
.slot_mapping.fill_(-1) .slot_mapping.fill_(-1)
if ubatch_slices is not None: if ubatch_slices is not None:
for ubid, (req_slice, token_slice) in enumerate(ubatch_slices): common_attn_metadata_list = split_attn_metadata(
assert token_slice.stop > token_slice.start ubatch_slices, common_attn_metadata)
for ubid, common_attn_metadata in enumerate(
common_attn_metadata_list):
assert common_attn_metadata.max_query_len == 1
attn_metadata_i = ( attn_metadata_i = (
self.attn_metadata_builders[kv_cache_group_id]. self.attn_metadata_builders[kv_cache_group_id].build(
build_slice(
req_slice=req_slice,
token_slice=token_slice,
max_query_len=max(tokens[req_slice]),
common_prefix_len=common_prefix_len, common_prefix_len=common_prefix_len,
common_attn_metadata=common_attn_metadata, common_attn_metadata=common_attn_metadata,
)) ubatch_id=ubid))
for layer_name in kv_cache_group_spec.layer_names: for layer_name in kv_cache_group_spec.layer_names:
assert type(attn_metadata) is list assert type(attn_metadata) is list
attn_metadata[ubid][layer_name] = attn_metadata_i attn_metadata[ubid][layer_name] = attn_metadata_i