biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Kernel][TPU][ragged-paged-attn] vLLM code change for PR#8896 (#15659)

Browse Source 
Signed-off-by: Yarong Mu <ymu@google.com>
This commit is contained in:
yarongmu-google
2025-03-28 21:13:15 -07:00
committed by GitHub
parent da461f3cbf
commit 7c1f760024
4 changed files with 37 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
vllm/v1/worker/tpu_model_runner.py
View File

@@ -861,12 +861,11 @@ class TPUModelRunner:
kv_cache_spec.num_kv_heads, kv_cache_spec.head_size)
dtype = kv_cache_spec.dtype
tpu_k_cache = torch.zeros(kv_cache_shape,
dtype=dtype,
device=self.device)
tpu_v_cache = torch.zeros_like(tpu_k_cache)
tpu_kv_cache = torch.zeros(kv_cache_shape,
dtype=dtype,
device=self.device)
kv_caches[layer_name] = (tpu_k_cache, tpu_v_cache)
kv_caches[layer_name] = tpu_kv_cache
else:
raise NotImplementedError
@@ -893,7 +892,7 @@ class ModelWrapperV1(nn.Module):
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
kv_caches: list[tuple[torch.Tensor, torch.Tensor]],
kv_caches: list[torch.Tensor],
inputs_embeds: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Executes the forward pass of the model.

8
vllm/v1/worker/tpu_worker.py
View File

@@ -136,10 +136,10 @@ class TPUWorker:
# Use an empty tensor instead of `None`` to force Dynamo to pass
# it by reference, rather by specializing on the value ``None``.
tpu_k_cache = torch.tensor([], dtype=dtype, device=self.device)
tpu_v_cache = torch.tensor([], dtype=dtype, device=self.device)
kv_caches[layer_name] = (tpu_k_cache, tpu_v_cache)
tpu_kv_cache = torch.tensor([],
dtype=dtype,
device=self.device)
kv_caches[layer_name] = tpu_kv_cache
else:
raise NotImplementedError