Skip to content

[ROCm] Resolve the indexer fp8 cache dtype once at import #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
amd-sriram wants to merge 1 commit into
base: rocm/fp8-hoist-base
Choose a base branch
from
Draft

[ROCm] Resolve the indexer fp8 cache dtype once at import #8

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 5 additions & 8 deletions vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,8 +109,7 @@ def indexer_k_quant_and_cache_triton(
# In real layout, we store the first portion as kv cache value
# and second portion as kv cache scale
kv_cache = kv_cache.view(num_blocks, -1)
fp8_dtype = current_platform.fp8_dtype()
kv_cache_value = kv_cache[:, : block_size * head_dim].view(fp8_dtype)
kv_cache_value = kv_cache[:, : block_size * head_dim].view(FP8_DTYPE)
kv_cache_scale = kv_cache[:, block_size * head_dim :].view(torch.float32)
head_tile_size = head_tile_size // kv_cache.element_size()
layout = "NORMAL" if block_size == 1 else "SHUFFLE"
Expand All @@ -128,7 +127,7 @@ def indexer_k_quant_and_cache_triton(
layout,
block_tile_size,
head_tile_size,
IS_FNUZ=current_platform.fp8_dtype() == torch.float8_e4m3fnuz,
IS_FNUZ=FP8_DTYPE == torch.float8_e4m3fnuz,
USE_UE8M0=scale_fmt == "ue8m0",
)

Expand Down Expand Up @@ -406,8 +405,7 @@ def cp_gather_indexer_k_quant_cache_triton(
num_blocks = k_cache.shape[0]
# we assume the kv cache already been split to 2 portion
k_cache = k_cache.view(num_blocks, -1)
fp8_dtype = current_platform.fp8_dtype()
k_cache_value = k_cache[:, : block_size * head_dim].view(fp8_dtype)
k_cache_value = k_cache[:, : block_size * head_dim].view(FP8_DTYPE)
k_cache_scale = k_cache[:, block_size * head_dim :].view(torch.float32)
grid = (num_tokens,)
k_fp8_scale = k_fp8_scale.view(torch.float32)
Expand Down Expand Up @@ -457,7 +455,6 @@ def fp8_paged_mqa_logits_torch(
):
from vllm.utils.math_utils import cdiv

fp8_dtype = current_platform.fp8_dtype()
batch_size, next_n, _, dim = q.size()
if next_n == 1:
block_size = kv_cache.shape[1]
Expand All @@ -481,7 +478,7 @@ def fp8_paged_mqa_logits_torch(
cache = kv_cache_flat[pages]
scale_offset = block_size * dim
cache_value = (
cache[..., :scale_offset].view(dtype=fp8_dtype).to(torch.float32)
cache[..., :scale_offset].view(dtype=FP8_DTYPE).to(torch.float32)
)
cache_scale = (
cache[..., scale_offset:].view(dtype=torch.float32).contiguous()
Expand All @@ -499,7 +496,7 @@ def fp8_paged_mqa_logits_torch(
kv_cache, scale = kv_cache[..., :dim], kv_cache[..., dim:]
scale = scale.contiguous().view(torch.float)
q = q.float()
kv_cache = kv_cache.view(fp8_dtype).float() * scale
kv_cache = kv_cache.view(FP8_DTYPE).float() * scale
num_block, block_size, _, dim = kv_cache.size()
logits = torch.full(
[batch_size * next_n, max_model_len],
Expand Down
Loading