forked from Karylab-cklius/vllm
[XPU] Enable multiple key kernels for sparse attention (#37888)
Signed-off-by: Xiaochang Wu <xiaochang.wu@intel.com> Signed-off-by: Wu, Xiaochang <xiaochang.wu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
Xiaochang Wu
co-authored by
Kunshang Ji
Kunshang Ji
parent
18a27cc9a3
commit
8c8b1825eb
@@ -2826,6 +2826,50 @@ def indexer_k_quant_and_cache(
|
||||
)
|
||||
|
||||
|
||||
def top_k_per_row_prefill(
|
||||
logits: torch.Tensor,
|
||||
cu_seqlen_ks: torch.Tensor,
|
||||
cu_seqlen_ke: torch.Tensor,
|
||||
raw_topk_indices: torch.Tensor,
|
||||
num_rows: int,
|
||||
stride0: int,
|
||||
stride1: int,
|
||||
topk_tokens: int,
|
||||
) -> None:
|
||||
torch.ops._C.top_k_per_row_prefill(
|
||||
logits,
|
||||
cu_seqlen_ks,
|
||||
cu_seqlen_ke,
|
||||
raw_topk_indices,
|
||||
num_rows,
|
||||
stride0,
|
||||
stride1,
|
||||
topk_tokens,
|
||||
)
|
||||
|
||||
|
||||
def top_k_per_row_decode(
|
||||
logits: torch.Tensor,
|
||||
next_n: int,
|
||||
seq_lens: torch.Tensor,
|
||||
raw_topk_indices: torch.Tensor,
|
||||
num_rows: int,
|
||||
stride0: int,
|
||||
stride1: int,
|
||||
topk_tokens: int,
|
||||
) -> None:
|
||||
torch.ops._C.top_k_per_row_decode(
|
||||
logits,
|
||||
next_n,
|
||||
seq_lens,
|
||||
raw_topk_indices,
|
||||
num_rows,
|
||||
stride0,
|
||||
stride1,
|
||||
topk_tokens,
|
||||
)
|
||||
|
||||
|
||||
def cp_gather_indexer_k_quant_cache(
|
||||
kv_cache: torch.Tensor,
|
||||
dst_k: torch.Tensor,
|
||||
|
||||
+82
-245
@@ -185,6 +185,76 @@ def _xpu_ops_deepseek_scaling_rope_fake(
|
||||
return query, key
|
||||
|
||||
|
||||
def _xpu_fp8_mqa_logits_impl(
|
||||
q: torch.Tensor,
|
||||
k_quant: torch.Tensor,
|
||||
k_scale: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
cu_seqlen_ks: torch.Tensor,
|
||||
cu_seqlen_ke: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return torch.ops._xpu_C.fp8_mqa_logits(
|
||||
q,
|
||||
k_quant,
|
||||
k_scale,
|
||||
weights,
|
||||
cu_seqlen_ks,
|
||||
cu_seqlen_ke,
|
||||
)
|
||||
|
||||
|
||||
def _xpu_fp8_mqa_logits_fake(
|
||||
q: torch.Tensor,
|
||||
k_quant: torch.Tensor,
|
||||
k_scale: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
cu_seqlen_ks: torch.Tensor,
|
||||
cu_seqlen_ke: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return torch.empty(
|
||||
(q.shape[0], k_quant.shape[0]),
|
||||
dtype=torch.float32,
|
||||
device=q.device,
|
||||
)
|
||||
|
||||
|
||||
def _xpu_fp8_paged_mqa_logits_impl(
|
||||
q: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
schedule_metadata: torch.Tensor,
|
||||
max_model_len: int,
|
||||
) -> torch.Tensor:
|
||||
return torch.ops._xpu_C.fp8_paged_mqa_logits(
|
||||
q,
|
||||
kv_cache,
|
||||
weights,
|
||||
context_lens,
|
||||
block_tables,
|
||||
schedule_metadata,
|
||||
max_model_len,
|
||||
)
|
||||
|
||||
|
||||
def _xpu_fp8_paged_mqa_logits_fake(
|
||||
q: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
block_tables: torch.Tensor,
|
||||
schedule_metadata: torch.Tensor,
|
||||
max_model_len: int,
|
||||
) -> torch.Tensor:
|
||||
batch_size, next_n = q.shape[:2]
|
||||
return torch.empty(
|
||||
(batch_size * next_n, max_model_len),
|
||||
dtype=torch.float32,
|
||||
device=q.device,
|
||||
)
|
||||
|
||||
|
||||
def _topk_topp_sample_impl(
|
||||
random_sampled: torch.Tensor,
|
||||
logits_to_return: torch.Tensor | None,
|
||||
@@ -438,251 +508,6 @@ class xpu_ops:
|
||||
)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def indexer_k_quant_and_cache(
|
||||
k: torch.Tensor,
|
||||
kv_cache: torch.Tensor,
|
||||
slot_mapping: torch.Tensor,
|
||||
quant_block_size: int,
|
||||
scale_fmt: str | None,
|
||||
) -> None:
|
||||
head_dim = k.shape[-1]
|
||||
k = k.view(-1, head_dim) # [total_tokens, head_dim]
|
||||
|
||||
def group_quant_torch(
|
||||
x: torch.Tensor,
|
||||
group_size: int,
|
||||
eps: float = 1e-10,
|
||||
dtype: torch.dtype | None = None,
|
||||
column_major_scales: bool = False,
|
||||
out_q: torch.Tensor | None = None,
|
||||
use_ue8m0: bool | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if use_ue8m0 is None:
|
||||
# Default fallback - could import is_deep_gemm_e8m0_used if needed
|
||||
use_ue8m0 = False
|
||||
|
||||
if dtype is None:
|
||||
dtype = current_platform.fp8_dtype()
|
||||
|
||||
# Validate inputs
|
||||
assert x.shape[-1] % group_size == 0, (
|
||||
f"Last dimension {x.shape[-1]} must be divisible by "
|
||||
f"group_size {group_size}"
|
||||
)
|
||||
assert x.stride(-1) == 1, "Input tensor groups must be contiguous"
|
||||
|
||||
# Prepare output tensor
|
||||
if out_q is None:
|
||||
x_q = torch.empty_like(x, dtype=dtype)
|
||||
else:
|
||||
assert out_q.shape == x.shape
|
||||
x_q = out_q
|
||||
|
||||
# Reshape input for group processing
|
||||
# Original shape: (..., last_dim)
|
||||
# Target shape: (..., num_groups, group_size)
|
||||
original_shape = x.shape
|
||||
num_groups = original_shape[-1] // group_size
|
||||
|
||||
# Reshape to separate groups
|
||||
group_shape = original_shape[:-1] + (num_groups, group_size)
|
||||
x_grouped = x.view(group_shape)
|
||||
|
||||
# Compute per-group absolute maximum values
|
||||
# Shape: (..., num_groups)
|
||||
abs_max = torch.amax(torch.abs(x_grouped), dim=-1, keepdim=False)
|
||||
abs_max = torch.maximum(
|
||||
abs_max, torch.tensor(eps, device=x.device, dtype=x.dtype)
|
||||
)
|
||||
|
||||
# Compute scales
|
||||
FP8_MAX = torch.finfo(dtype).max
|
||||
FP8_MIN = torch.finfo(dtype).min
|
||||
scale_raw = abs_max / FP8_MAX
|
||||
|
||||
if use_ue8m0:
|
||||
# For UE8M0 format, scales must be powers of 2
|
||||
scales = torch.pow(2.0, torch.ceil(torch.log2(scale_raw)))
|
||||
else:
|
||||
scales = scale_raw
|
||||
|
||||
# Expand scales for broadcasting with grouped data
|
||||
# Shape: (..., num_groups, 1)
|
||||
scales_expanded = scales.unsqueeze(-1)
|
||||
|
||||
# Quantize the grouped data
|
||||
x_scaled = x_grouped / scales_expanded
|
||||
x_clamped = torch.clamp(x_scaled, FP8_MIN, FP8_MAX)
|
||||
x_quantized = x_clamped.to(dtype)
|
||||
|
||||
# Reshape back to original shape
|
||||
x_q.copy_(x_quantized.view(original_shape))
|
||||
|
||||
# Prepare scales tensor in requested format
|
||||
if column_major_scales:
|
||||
# Column-major: (num_groups,) + batch_dims
|
||||
# Transpose the scales to put group dimension first
|
||||
scales_shape = (num_groups,) + original_shape[:-1]
|
||||
x_s = scales.permute(-1, *range(len(original_shape) - 1))
|
||||
x_s = x_s.contiguous().view(scales_shape)
|
||||
else:
|
||||
# Row-major: batch_dims + (num_groups,)
|
||||
x_s = scales.contiguous()
|
||||
|
||||
# Ensure scales are float32
|
||||
return x_q, x_s.float()
|
||||
|
||||
k_fp8, k_scale = group_quant_torch(
|
||||
k,
|
||||
group_size=quant_block_size,
|
||||
column_major_scales=False,
|
||||
use_ue8m0=(scale_fmt == "ue8m0"),
|
||||
)
|
||||
|
||||
k_fp8_bytes = k_fp8.view(-1, head_dim).view(torch.uint8)
|
||||
scale_bytes = k_scale.view(torch.uint8).view(-1, 4)
|
||||
k = torch.cat(
|
||||
[k_fp8_bytes, scale_bytes], dim=-1
|
||||
) # [total_tokens, head_dim + 4]
|
||||
|
||||
slot_mapping = slot_mapping.flatten()
|
||||
# kv_cache: [num_block, block_size, head_dim + 4]
|
||||
kv_cache.view(-1, kv_cache.shape[-1]).index_copy_(0, slot_mapping, k)
|
||||
|
||||
@staticmethod
|
||||
def cp_gather_indexer_k_quant_cache(
|
||||
kv_cache: torch.Tensor,
|
||||
dst_k: torch.Tensor,
|
||||
dst_scale: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
cu_seq_lens: torch.Tensor,
|
||||
) -> None:
|
||||
"""
|
||||
Args:
|
||||
kv_cache: [num_blocks, block_size, cache_stride] - quantized KV cache
|
||||
Layout per block: [k_values, scale_values]
|
||||
- k_values: [block_size * head_dim]
|
||||
- scale_values: [block_size * head_dim * 4 / quant_block_size]
|
||||
dst_k: [num_tokens, head_dim] - output tensor for K values
|
||||
dst_scale: [num_tokens, head_dim / quant_block_size * 4]
|
||||
- output tensor for scale values
|
||||
block_table: [batch_size, num_blocks] - block table for indexing
|
||||
cu_seq_lens: [batch_size + 1] - cumulative sequence lengths
|
||||
"""
|
||||
batch_size = block_table.size(0)
|
||||
num_tokens = dst_k.size(0)
|
||||
head_dim = dst_k.size(1)
|
||||
cache_block_size = kv_cache.size(1)
|
||||
quant_block_size = head_dim * 4 // dst_scale.size(1)
|
||||
|
||||
# For each token, find which batch it belongs to using searchsorted
|
||||
token_indices = torch.arange(num_tokens, device=dst_k.device) + 1
|
||||
# cu_seq_lens is [batch_size + 1], we need to find which interval each
|
||||
# token belongs to
|
||||
batch_indices = torch.searchsorted(cu_seq_lens, token_indices) - 1
|
||||
batch_indices = torch.clamp(batch_indices, 0, batch_size - 1)
|
||||
|
||||
# Calculate the in-batch sequence index for each token
|
||||
inbatch_seq_indices = token_indices - cu_seq_lens[batch_indices]
|
||||
|
||||
# Find which block each token belongs to
|
||||
block_indices_in_table = inbatch_seq_indices // cache_block_size
|
||||
physical_block_indices = block_table[batch_indices, block_indices_in_table]
|
||||
|
||||
# Calculate the offset within each block
|
||||
inblock_offsets = (inbatch_seq_indices - 1) % cache_block_size
|
||||
|
||||
# Calculate strides
|
||||
block_stride = kv_cache.stride(0) # stride for each block
|
||||
|
||||
# Flatten kv_cache for easier indexing
|
||||
kv_cache_flat = kv_cache.view(-1)
|
||||
|
||||
# Calculate source offset for K values for all tokens (vectorized)
|
||||
src_block_offsets = physical_block_indices * block_stride
|
||||
src_k_offsets = src_block_offsets + inblock_offsets * head_dim
|
||||
|
||||
# Gather K values using advanced indexing
|
||||
# Create indices for all elements we need to gather
|
||||
k_indices = src_k_offsets.unsqueeze(1) + torch.arange(
|
||||
head_dim, device=dst_k.device
|
||||
)
|
||||
dst_k[:] = kv_cache_flat[k_indices]
|
||||
|
||||
# Calculate source offset for scale values (vectorized)
|
||||
# Scales are stored after all K values for each block
|
||||
scale_size = head_dim * 4 // quant_block_size
|
||||
src_scale_offsets = src_block_offsets + head_dim + inblock_offsets * scale_size
|
||||
|
||||
# Gather scale values
|
||||
scale_indices = src_scale_offsets.unsqueeze(1) + torch.arange(
|
||||
scale_size, device=dst_scale.device
|
||||
)
|
||||
dst_scale[:] = kv_cache_flat[scale_indices]
|
||||
|
||||
@staticmethod
|
||||
def top_k_per_row_prefill(
|
||||
logits: torch.Tensor,
|
||||
cu_seqlen_ks: torch.Tensor,
|
||||
cu_seqlen_ke: torch.Tensor,
|
||||
raw_topk_indices: torch.Tensor,
|
||||
num_rows: int,
|
||||
stride0: int,
|
||||
strdide1: int,
|
||||
topk_tokens: int,
|
||||
) -> torch.Tensor:
|
||||
real_topk = min(topk_tokens, logits.shape[-1])
|
||||
topk_indices = logits.topk(real_topk, dim=-1)[1].to(torch.int32)
|
||||
topk_indices -= cu_seqlen_ks[:, None]
|
||||
mask_lo = topk_indices >= 0
|
||||
mask_hi = topk_indices - (cu_seqlen_ke - cu_seqlen_ks)[:, None] < 0
|
||||
mask = torch.full_like(
|
||||
topk_indices, False, dtype=torch.bool, device=topk_indices.device
|
||||
)
|
||||
mask = mask_lo & mask_hi
|
||||
topk_indices.masked_fill_(~mask, -1)
|
||||
raw_topk_indices[: topk_indices.shape[0], : topk_indices.shape[1]] = (
|
||||
topk_indices
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def top_k_per_row_decode(
|
||||
logits: torch.Tensor,
|
||||
next_n: int,
|
||||
seq_lens: torch.Tensor,
|
||||
raw_topk_indices: torch.Tensor,
|
||||
num_rows: int,
|
||||
stride0: int,
|
||||
stride1: int,
|
||||
topk_tokens: int,
|
||||
) -> torch.Tensor:
|
||||
device = logits.device
|
||||
batch_size = seq_lens.size(0)
|
||||
# padded query len
|
||||
padded_num_tokens = batch_size * next_n
|
||||
positions = (
|
||||
torch.arange(logits.shape[-1], device=device)
|
||||
.unsqueeze(0)
|
||||
.expand(batch_size * next_n, -1)
|
||||
)
|
||||
row_indices = torch.arange(padded_num_tokens, device=device) // next_n
|
||||
next_n_offset = torch.arange(padded_num_tokens, device=device) % next_n
|
||||
index_end_pos = (seq_lens[row_indices] - next_n + next_n_offset).unsqueeze(1)
|
||||
# index_end_pos: [B * N, 1]
|
||||
mask = positions <= index_end_pos
|
||||
# mask: [B * N, L]
|
||||
logits = logits.masked_fill(~mask, float("-inf"))
|
||||
real_topk = min(topk_tokens, logits.shape[-1])
|
||||
topk_indices = logits.topk(real_topk, dim=-1)[1].to(torch.int32) # [B * N, K]
|
||||
# ensure we don't set indices for the top k
|
||||
# that is out of range(masked already)
|
||||
# this will happen if context length is shorter than K
|
||||
topk_indices[topk_indices > index_end_pos] = -1
|
||||
raw_topk_indices[: topk_indices.shape[0], : topk_indices.shape[1]] = (
|
||||
topk_indices
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def register_ops_once() -> None:
|
||||
global _OPS_REGISTERED
|
||||
@@ -708,6 +533,18 @@ class xpu_ops:
|
||||
fake_impl=_xpu_mxfp4_quantize_fake,
|
||||
)
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="xpu_fp8_mqa_logits",
|
||||
op_func=_xpu_fp8_mqa_logits_impl,
|
||||
fake_impl=_xpu_fp8_mqa_logits_fake,
|
||||
)
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="xpu_fp8_paged_mqa_logits",
|
||||
op_func=_xpu_fp8_paged_mqa_logits_impl,
|
||||
fake_impl=_xpu_fp8_paged_mqa_logits_fake,
|
||||
)
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="gdn_attention_core_xpu",
|
||||
op_func=_gdn_attention_core_xpu_impl,
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
import torch
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
|
||||
from vllm.forward_context import get_forward_context
|
||||
@@ -28,11 +29,6 @@ from vllm.v1.attention.backends.mla.indexer import (
|
||||
from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
|
||||
from vllm.v1.worker.workspace import current_workspace_manager
|
||||
|
||||
if current_platform.is_cuda_alike():
|
||||
from vllm import _custom_ops as ops
|
||||
elif current_platform.is_xpu():
|
||||
from vllm._xpu_ops import xpu_ops
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
|
||||
@@ -222,42 +218,42 @@ def sparse_attn_indexer(
|
||||
q_slice_cast = q_slice
|
||||
k_quant_cast = k_quant
|
||||
k_scale_cast = k_scale.view(torch.float32).squeeze(-1)
|
||||
logits = fp8_fp4_mqa_logits(
|
||||
(q_slice_cast, q_scale_slice),
|
||||
(k_quant_cast, k_scale_cast),
|
||||
weights[chunk.token_start : chunk.token_end],
|
||||
chunk.cu_seqlen_ks,
|
||||
chunk.cu_seqlen_ke,
|
||||
clean_logits=False,
|
||||
)
|
||||
if current_platform.is_xpu():
|
||||
if q_scale_slice is not None:
|
||||
raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q")
|
||||
logits = torch.ops.vllm.xpu_fp8_mqa_logits(
|
||||
q_slice_cast,
|
||||
k_quant_cast,
|
||||
k_scale_cast,
|
||||
weights[chunk.token_start : chunk.token_end],
|
||||
chunk.cu_seqlen_ks,
|
||||
chunk.cu_seqlen_ke,
|
||||
)
|
||||
else:
|
||||
logits = fp8_fp4_mqa_logits(
|
||||
(q_slice_cast, q_scale_slice),
|
||||
(k_quant_cast, k_scale_cast),
|
||||
weights[chunk.token_start : chunk.token_end],
|
||||
chunk.cu_seqlen_ks,
|
||||
chunk.cu_seqlen_ke,
|
||||
clean_logits=False,
|
||||
)
|
||||
num_rows = logits.shape[0]
|
||||
|
||||
topk_indices = topk_indices_buffer[
|
||||
chunk.token_start : chunk.token_end, :topk_tokens
|
||||
]
|
||||
|
||||
if current_platform.is_xpu():
|
||||
xpu_ops.top_k_per_row_prefill( # type: ignore[attr-defined]
|
||||
logits,
|
||||
chunk.cu_seqlen_ks,
|
||||
chunk.cu_seqlen_ke,
|
||||
topk_indices,
|
||||
num_rows,
|
||||
logits.stride(0),
|
||||
logits.stride(1),
|
||||
topk_tokens,
|
||||
)
|
||||
else:
|
||||
torch.ops._C.top_k_per_row_prefill(
|
||||
logits,
|
||||
chunk.cu_seqlen_ks,
|
||||
chunk.cu_seqlen_ke,
|
||||
topk_indices,
|
||||
num_rows,
|
||||
logits.stride(0),
|
||||
logits.stride(1),
|
||||
topk_tokens,
|
||||
)
|
||||
ops.top_k_per_row_prefill(
|
||||
logits,
|
||||
chunk.cu_seqlen_ks,
|
||||
chunk.cu_seqlen_ke,
|
||||
topk_indices,
|
||||
num_rows,
|
||||
logits.stride(0),
|
||||
logits.stride(1),
|
||||
topk_tokens,
|
||||
)
|
||||
|
||||
if has_decode:
|
||||
decode_metadata = attn_metadata_narrowed.decode
|
||||
@@ -309,16 +305,32 @@ def sparse_attn_indexer(
|
||||
if use_fp4_cache
|
||||
else padded_q_quant_decode_tokens
|
||||
)
|
||||
logits = fp8_fp4_paged_mqa_logits(
|
||||
(padded_q_quant_cast, padded_q_scale),
|
||||
kv_cache,
|
||||
weights[:num_padded_tokens],
|
||||
seq_lens,
|
||||
decode_metadata.block_table,
|
||||
decode_metadata.schedule_metadata,
|
||||
max_model_len=max_model_len,
|
||||
clean_logits=False,
|
||||
)
|
||||
if current_platform.is_xpu():
|
||||
if padded_q_scale is not None:
|
||||
raise RuntimeError("XPU fp8_paged_mqa_logits does not support FP4 Q")
|
||||
seq_lens_xpu = (
|
||||
seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens
|
||||
)
|
||||
logits = torch.ops.vllm.xpu_fp8_paged_mqa_logits(
|
||||
padded_q_quant_cast,
|
||||
kv_cache,
|
||||
weights[:num_padded_tokens],
|
||||
seq_lens_xpu,
|
||||
decode_metadata.block_table,
|
||||
decode_metadata.schedule_metadata,
|
||||
max_model_len,
|
||||
)
|
||||
else:
|
||||
logits = fp8_fp4_paged_mqa_logits(
|
||||
(padded_q_quant_cast, padded_q_scale),
|
||||
kv_cache,
|
||||
weights[:num_padded_tokens],
|
||||
seq_lens,
|
||||
decode_metadata.block_table,
|
||||
decode_metadata.schedule_metadata,
|
||||
max_model_len=max_model_len,
|
||||
clean_logits=False,
|
||||
)
|
||||
num_rows = logits.shape[0]
|
||||
topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
|
||||
|
||||
@@ -336,28 +348,16 @@ def sparse_attn_indexer(
|
||||
attn_metadata_narrowed.max_seq_len,
|
||||
)
|
||||
else:
|
||||
if current_platform.is_xpu():
|
||||
xpu_ops.top_k_per_row_decode( # type: ignore[attr-defined]
|
||||
logits,
|
||||
next_n,
|
||||
seq_lens,
|
||||
topk_indices,
|
||||
num_rows,
|
||||
logits.stride(0),
|
||||
logits.stride(1),
|
||||
topk_tokens,
|
||||
)
|
||||
else:
|
||||
torch.ops._C.top_k_per_row_decode(
|
||||
logits,
|
||||
next_n,
|
||||
seq_lens,
|
||||
topk_indices,
|
||||
num_rows,
|
||||
logits.stride(0),
|
||||
logits.stride(1),
|
||||
topk_tokens,
|
||||
)
|
||||
ops.top_k_per_row_decode(
|
||||
logits,
|
||||
next_n,
|
||||
seq_lens,
|
||||
topk_indices,
|
||||
num_rows,
|
||||
logits.stride(0),
|
||||
logits.stride(1),
|
||||
topk_tokens,
|
||||
)
|
||||
|
||||
if decode_metadata.requires_padding:
|
||||
# if padded, we need to unpack
|
||||
@@ -494,6 +494,15 @@ class SparseAttnIndexer(CustomOp):
|
||||
self.use_fp4_cache,
|
||||
)
|
||||
|
||||
def forward_xpu(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
q_fp8: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
):
|
||||
return self.forward_cuda(hidden_states, q_fp8, k, weights)
|
||||
|
||||
def forward_hip(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
|
||||
Reference in New Issue
Block a user