wehub-resource-sync
1085 lines
35 KiB
Python
1085 lines
35 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from collections.abc import Callable
|
|
from typing import TYPE_CHECKING
|
|
|
|
import torch
|
|
from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.platforms import current_platform
|
|
from vllm.triton_utils import tl, triton
|
|
from vllm.utils.torch_utils import direct_register_custom_op
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
if TYPE_CHECKING:
|
|
|
|
def register_fake(fn):
|
|
return lambda name: fn
|
|
else:
|
|
try:
|
|
from torch.library import register_fake
|
|
except ImportError:
|
|
from torch.library import impl_abstract as register_fake
|
|
|
|
if hasattr(torch.ops._xpu_C, "fp8_gemm"):
|
|
|
|
@register_fake("_xpu_C::fp8_gemm")
|
|
def _fp8_gemm_fake(
|
|
q_input: torch.Tensor,
|
|
q_weight: torch.Tensor,
|
|
out_dtype: torch.dtype,
|
|
input_scales: torch.Tensor,
|
|
weight_scale: torch.Tensor,
|
|
bias: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
input_2d = q_input.view(-1, q_input.shape[-1])
|
|
M = input_2d.size(0)
|
|
N = q_weight.size(1)
|
|
return torch.empty((M, N), dtype=out_dtype, device=q_input.device)
|
|
|
|
|
|
if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"):
|
|
|
|
@register_fake("_xpu_C::fp8_gemm_w8a16")
|
|
def _fp8_gemm_w8a16_fake(
|
|
input: torch.Tensor,
|
|
q_weight: torch.Tensor,
|
|
weight_scale: torch.Tensor,
|
|
bias: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
input_2d = input.view(-1, input.shape[-1])
|
|
M = input_2d.size(0)
|
|
N = q_weight.size(1)
|
|
return torch.empty((M, N), dtype=input.dtype, device=input.device)
|
|
|
|
|
|
if hasattr(torch.ops._xpu_C, "int4_gemm_w4a8"):
|
|
|
|
@register_fake("_xpu_C::int4_gemm_w4a8")
|
|
def _int4_gemm_w4a8_fake(
|
|
input: torch.Tensor,
|
|
input_scales: torch.Tensor,
|
|
input_zero_points: torch.Tensor,
|
|
q_weight: torch.Tensor,
|
|
weight_scale: torch.Tensor,
|
|
weight_zp: torch.Tensor,
|
|
group_size: int,
|
|
g_idx: torch.Tensor | None = None,
|
|
bias: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
input_2d = input.view(-1, input.shape[-1])
|
|
M = input_2d.size(0)
|
|
N = q_weight.size(1)
|
|
return torch.empty((M, N), dtype=torch.float16, device=input.device)
|
|
|
|
|
|
if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"):
|
|
|
|
@register_fake("_xpu_C::int4_gemm_w4a16")
|
|
def _int4_gemm_w4a16_fake(
|
|
input: torch.Tensor,
|
|
q_weight: torch.Tensor,
|
|
bias: torch.Tensor | None,
|
|
weight_scale: torch.Tensor,
|
|
qzeros: torch.Tensor,
|
|
group_size: int,
|
|
group_idx: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
input_2d = input.view(-1, input.shape[-1])
|
|
M = input_2d.size(0)
|
|
N = q_weight.size(1)
|
|
return torch.empty((M, N), dtype=input.dtype, device=input.device)
|
|
|
|
|
|
def _gdn_attention_core_xpu_impl(
|
|
core_attn_out: torch.Tensor,
|
|
z: torch.Tensor,
|
|
projected_states_qkvz: torch.Tensor,
|
|
projected_states_ba: torch.Tensor,
|
|
layer_name: str,
|
|
) -> None:
|
|
"""Custom op wrapping the XPU SYCL GDN kernel for torch.compile."""
|
|
from vllm.forward_context import get_forward_context
|
|
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
|
|
|
|
forward_context = get_forward_context()
|
|
self = forward_context.no_compile_layers[layer_name]
|
|
attn_metadata_raw = forward_context.attn_metadata
|
|
|
|
if attn_metadata_raw is None:
|
|
return
|
|
|
|
assert isinstance(attn_metadata_raw, dict)
|
|
attn_metadata = attn_metadata_raw[self.prefix]
|
|
assert isinstance(attn_metadata, GDNAttentionMetadata)
|
|
|
|
num_actual_tokens = attn_metadata.num_actual_tokens
|
|
num_accepted_tokens = attn_metadata.num_accepted_tokens
|
|
|
|
num_prefills = attn_metadata.num_prefills
|
|
num_decodes = attn_metadata.num_decodes
|
|
num_spec_decodes = attn_metadata.num_spec_decodes
|
|
|
|
has_initial_state = attn_metadata.has_initial_state
|
|
|
|
non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
|
|
non_spec_token_indx = attn_metadata.non_spec_token_indx
|
|
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
|
|
non_spec_state_indices_tensor = (
|
|
non_spec_state_indices_tensor.contiguous()
|
|
if non_spec_state_indices_tensor is not None
|
|
else None
|
|
)
|
|
|
|
spec_query_start_loc = attn_metadata.spec_query_start_loc
|
|
spec_token_indx = attn_metadata.spec_token_indx
|
|
spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501
|
|
|
|
spec_sequence_masks = attn_metadata.spec_sequence_masks
|
|
if spec_sequence_masks is not None:
|
|
if non_spec_token_indx is not None:
|
|
non_spec_token_indx = non_spec_token_indx.to(torch.int32)
|
|
if spec_token_indx is not None:
|
|
spec_token_indx = spec_token_indx.to(torch.int32)
|
|
|
|
conv_weights = self.conv1d.weight.view(
|
|
self.conv1d.weight.size(0), self.conv1d.weight.size(2)
|
|
)
|
|
|
|
torch.ops._xpu_C.gdn_attention(
|
|
core_attn_out,
|
|
z,
|
|
projected_states_qkvz,
|
|
projected_states_ba,
|
|
self.num_k_heads,
|
|
self.num_v_heads,
|
|
self.head_k_dim,
|
|
self.head_v_dim,
|
|
conv_state=self.kv_cache[0],
|
|
ssm_state=self.kv_cache[1],
|
|
conv_weights=conv_weights,
|
|
conv_bias=self.conv1d.bias,
|
|
activation=self.activation,
|
|
A_log=self.A_log,
|
|
dt_bias=self.dt_bias,
|
|
num_prefills=num_prefills, # type: ignore[attr-defined]
|
|
num_decodes=num_decodes, # type: ignore[attr-defined]
|
|
num_spec_decodes=num_spec_decodes, # type: ignore[attr-defined]
|
|
has_initial_state=has_initial_state, # type: ignore[attr-defined]
|
|
non_spec_query_start_loc=non_spec_query_start_loc, # type: ignore[attr-defined]
|
|
non_spec_token_indx=non_spec_token_indx, # type: ignore[attr-defined]
|
|
non_spec_state_indices_tensor=non_spec_state_indices_tensor, # type: ignore[attr-defined]
|
|
spec_query_start_loc=spec_query_start_loc, # type: ignore[attr-defined]
|
|
spec_token_indx=spec_token_indx, # type: ignore[attr-defined]
|
|
spec_state_indices_tensor=spec_state_indices_tensor,
|
|
num_accepted_tokens=num_accepted_tokens, # type: ignore[attr-defined]
|
|
num_actual_tokens=num_actual_tokens, # type: ignore[attr-defined]
|
|
tp_size=self.tp_size,
|
|
reorder_input=not self.gqa_interleaved_layout,
|
|
)
|
|
|
|
|
|
def _gdn_attention_core_xpu_fake(
|
|
core_attn_out: torch.Tensor,
|
|
z: torch.Tensor,
|
|
projected_states_qkvz: torch.Tensor,
|
|
projected_states_ba: torch.Tensor,
|
|
layer_name: str,
|
|
) -> None:
|
|
return
|
|
|
|
|
|
def _xpu_ops_deepseek_scaling_rope_impl(
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor | None,
|
|
offsets: torch.Tensor | None,
|
|
cos_sin_cache: torch.Tensor | None,
|
|
rotary_dim: int,
|
|
is_neox_style: bool,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
assert key is not None
|
|
return torch.ops._xpu_C.deepseek_scaling_rope(
|
|
positions, query, key, offsets, cos_sin_cache, rotary_dim, is_neox_style
|
|
)
|
|
|
|
|
|
def _xpu_ops_deepseek_scaling_rope_fake(
|
|
positions: torch.Tensor,
|
|
query: torch.Tensor,
|
|
key: torch.Tensor | None,
|
|
offsets: torch.Tensor | None,
|
|
cos_sin_cache: torch.Tensor | None,
|
|
rotary_dim: int,
|
|
is_neox_style: bool,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
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,
|
|
logits: torch.Tensor,
|
|
k: torch.Tensor | None,
|
|
p: torch.Tensor | None,
|
|
logprobs_mode: str,
|
|
seeds: torch.Tensor | None,
|
|
lambda_: float = 1.0,
|
|
) -> None:
|
|
torch.ops._xpu_C.topk_topp_sampler(
|
|
random_sampled, logits_to_return, logits, k, p, logprobs_mode, seeds, lambda_
|
|
)
|
|
return
|
|
|
|
|
|
def _topk_topp_sample_fake(
|
|
random_sampled: torch.Tensor,
|
|
logits_to_return: torch.Tensor | None,
|
|
logits: torch.Tensor,
|
|
k: torch.Tensor | None,
|
|
p: torch.Tensor | None,
|
|
logprobs_mode: str,
|
|
seeds: torch.Tensor | None,
|
|
lambda_: float = 1.0,
|
|
) -> None:
|
|
return
|
|
|
|
|
|
def _xpu_mxfp8_quantize_impl(
|
|
x: torch.Tensor, dtype: torch.dtype | None = None
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
MXFP8_BLOCK_SIZE = 32
|
|
assert x.shape[-1] % MXFP8_BLOCK_SIZE == 0
|
|
if dtype is not None:
|
|
assert dtype in (torch.float8_e4m3fn, torch.float8_e5m2), (
|
|
f"Unsupported dtype for xpu_mxfp8_quantize: {dtype}. "
|
|
f"Expected torch.float8_e4m3fn or torch.float8_e5m2."
|
|
)
|
|
else:
|
|
dtype = current_platform.fp8_dtype()
|
|
|
|
finfo = torch.finfo(dtype)
|
|
fp8_min = finfo.min
|
|
fp8_max = finfo.max
|
|
eps = 1e-10
|
|
|
|
x_q = torch.empty_like(x, device=x.device, dtype=dtype)
|
|
shape = x.shape[:-1] + (x.shape[-1] // MXFP8_BLOCK_SIZE,)
|
|
x_s = torch.empty(shape, device=x.device, dtype=torch.float32)
|
|
torch.ops._C.per_token_group_fp8_quant(
|
|
x,
|
|
x_q,
|
|
x_s,
|
|
MXFP8_BLOCK_SIZE,
|
|
eps,
|
|
fp8_min,
|
|
fp8_max,
|
|
True,
|
|
False,
|
|
False, # dummy_is_scale_transposed, dummy_is_tma_aligned
|
|
)
|
|
x_s = x_s.to(torch.float8_e8m0fnu)
|
|
return x_q, x_s
|
|
|
|
|
|
def _xpu_mxfp8_quantize_fake(
|
|
x: torch.Tensor, dtype: torch.dtype | None = None
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if dtype is None:
|
|
dtype = current_platform.fp8_dtype()
|
|
|
|
MXFP8_BLOCK_SIZE = 32
|
|
|
|
shape = x.shape[:-1] + (x.shape[-1] // MXFP8_BLOCK_SIZE,)
|
|
x_s = torch.zeros(shape, device=x.device, dtype=torch.float32)
|
|
|
|
return x.to(dtype), x_s.to(torch.float8_e8m0fnu)
|
|
|
|
|
|
def _xpu_mxfp4_quantize_impl(
|
|
x: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
MXFP4_BLOCK_SIZE = 32
|
|
eps = 1e-10
|
|
assert x.ndim == 2, "input must be 2-D"
|
|
assert x.shape[-1] % MXFP4_BLOCK_SIZE == 0, (
|
|
f"last dimension {x.shape[-1]} must be divisible by group_size "
|
|
f"{MXFP4_BLOCK_SIZE}"
|
|
)
|
|
assert x.is_contiguous(), "input groups must be contiguous"
|
|
|
|
M, N = x.shape
|
|
|
|
# Packed FP4 output: two nibbles per byte
|
|
x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8)
|
|
x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32)
|
|
|
|
torch.ops._C.per_token_group_quant_mxfp4(x, x_q, x_s, MXFP4_BLOCK_SIZE, eps)
|
|
|
|
x_q = x_q.view(torch.float4_e2m1fn_x2)
|
|
x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format)
|
|
return x_q, x_s
|
|
|
|
|
|
def _xpu_mxfp4_quantize_fake(
|
|
x: torch.Tensor,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
MXFP4_BLOCK_SIZE = 32
|
|
M, N = x.shape
|
|
|
|
# Packed FP4 output: two nibbles per byte
|
|
x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8)
|
|
x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32)
|
|
|
|
x_q = x_q.view(torch.float4_e2m1fn_x2)
|
|
x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format)
|
|
return x_q, x_s
|
|
|
|
|
|
@triton.jit
|
|
def _softplus(x):
|
|
return tl.where(x <= 20.0, tl.math.log(tl.math.exp(x) + 1.0), x)
|
|
|
|
|
|
@triton.jit
|
|
def _selective_scan_fwd_kernel(
|
|
# Pointers to input tensors
|
|
u_ptr,
|
|
delta_ptr,
|
|
A_ptr,
|
|
B_ptr,
|
|
C_ptr,
|
|
D_ptr,
|
|
z_ptr,
|
|
delta_bias_ptr,
|
|
# Pointers to output tensors (out aliases delta, out_z aliases z)
|
|
out_ptr,
|
|
out_z_ptr,
|
|
# SSM states
|
|
ssm_states_ptr,
|
|
# Optional pointers
|
|
query_start_loc_ptr,
|
|
cache_indices_ptr,
|
|
has_initial_state_ptr,
|
|
# APC pointers
|
|
block_idx_first_ptr,
|
|
block_idx_last_ptr,
|
|
initial_state_idx_ptr,
|
|
cu_chunk_seqlen_ptr,
|
|
last_chunk_indices_ptr,
|
|
# Dimensions
|
|
batch: tl.int32,
|
|
dim: tl.int32,
|
|
seqlen: tl.int32,
|
|
dstate: tl.int32,
|
|
n_groups: tl.int32,
|
|
dim_ngroups_ratio: tl.int32,
|
|
# Strides for u (and out, since out = delta which has same layout)
|
|
u_batch_stride: tl.int64,
|
|
u_d_stride: tl.int64,
|
|
# Strides for delta
|
|
delta_batch_stride: tl.int64,
|
|
delta_d_stride: tl.int64,
|
|
# Strides for A
|
|
A_d_stride: tl.int64,
|
|
A_dstate_stride: tl.int64,
|
|
# Strides for B
|
|
B_batch_stride: tl.int64,
|
|
B_group_stride: tl.int64,
|
|
B_dstate_stride: tl.int64,
|
|
# Strides for C
|
|
C_batch_stride: tl.int64,
|
|
C_group_stride: tl.int64,
|
|
C_dstate_stride: tl.int64,
|
|
# Strides for z
|
|
z_batch_stride: tl.int64,
|
|
z_d_stride: tl.int64,
|
|
# Strides for out
|
|
out_batch_stride: tl.int64,
|
|
out_d_stride: tl.int64,
|
|
# Strides for out_z
|
|
out_z_batch_stride: tl.int64,
|
|
out_z_d_stride: tl.int64,
|
|
# Strides for ssm_states
|
|
ssm_batch_stride: tl.int64,
|
|
ssm_dim_stride: tl.int64,
|
|
ssm_dstate_stride: tl.int64,
|
|
# Cache strides
|
|
cache_indices_stride: tl.int64,
|
|
# Scalar params
|
|
null_block_id: tl.int64,
|
|
block_size: tl.int32,
|
|
# Compile-time constants
|
|
delta_softplus: tl.constexpr,
|
|
HAS_D: tl.constexpr,
|
|
HAS_Z: tl.constexpr,
|
|
HAS_DELTA_BIAS: tl.constexpr,
|
|
IS_VARLEN: tl.constexpr,
|
|
HAS_CACHE_INDICES: tl.constexpr,
|
|
CACHE_ENABLED: tl.constexpr,
|
|
BLOCK_DSTATE: tl.constexpr,
|
|
):
|
|
batch_idx = tl.program_id(0)
|
|
dim_idx = tl.program_id(1)
|
|
group_idx = dim_idx // dim_ngroups_ratio
|
|
|
|
# Determine sequence boundaries
|
|
if IS_VARLEN:
|
|
seq_start = tl.load(query_start_loc_ptr + batch_idx).to(tl.int32)
|
|
seq_end = tl.load(query_start_loc_ptr + batch_idx + 1).to(tl.int32)
|
|
actual_seqlen = seq_end - seq_start
|
|
else:
|
|
seq_start = 0
|
|
actual_seqlen = seqlen
|
|
|
|
# Determine cache index for ssm_states
|
|
if CACHE_ENABLED:
|
|
init_state_idx = tl.load(initial_state_idx_ptr + batch_idx).to(tl.int32)
|
|
load_cache_slot = tl.load(
|
|
cache_indices_ptr + batch_idx * cache_indices_stride + init_state_idx
|
|
).to(tl.int64)
|
|
if load_cache_slot == null_block_id:
|
|
return
|
|
elif HAS_CACHE_INDICES:
|
|
cache_index = tl.load(cache_indices_ptr + batch_idx).to(tl.int64)
|
|
if cache_index == null_block_id:
|
|
return
|
|
load_cache_slot = cache_index
|
|
else:
|
|
load_cache_slot = batch_idx.to(tl.int64)
|
|
|
|
# Load D value
|
|
D_val = 0.0
|
|
if HAS_D:
|
|
D_val = tl.load(D_ptr + dim_idx).to(tl.float32)
|
|
|
|
# Load delta_bias value
|
|
delta_bias_val = 0.0
|
|
if HAS_DELTA_BIAS:
|
|
delta_bias_val = tl.load(delta_bias_ptr + dim_idx).to(tl.float32)
|
|
|
|
# Load A values for this dim - shape (dstate,)
|
|
dstate_offs = tl.arange(0, BLOCK_DSTATE)
|
|
dstate_mask = dstate_offs < dstate
|
|
A_vals = tl.load(
|
|
A_ptr + dim_idx * A_d_stride + dstate_offs * A_dstate_stride,
|
|
mask=dstate_mask,
|
|
other=0.0,
|
|
).to(tl.float32)
|
|
|
|
# Initialize state vector
|
|
state = tl.zeros((BLOCK_DSTATE,), dtype=tl.float32)
|
|
|
|
# Load initial state if available
|
|
has_init = False
|
|
if has_initial_state_ptr is not None:
|
|
has_init = tl.load(has_initial_state_ptr + batch_idx)
|
|
if has_init:
|
|
state = tl.load(
|
|
ssm_states_ptr
|
|
+ load_cache_slot * ssm_batch_stride
|
|
+ dim_idx * ssm_dim_stride
|
|
+ dstate_offs * ssm_dstate_stride,
|
|
mask=dstate_mask,
|
|
other=0.0,
|
|
).to(tl.float32)
|
|
|
|
# Compute base addresses for u and delta
|
|
if IS_VARLEN:
|
|
u_base = u_ptr + dim_idx * u_d_stride + seq_start * u_batch_stride
|
|
delta_base = (
|
|
delta_ptr + dim_idx * delta_d_stride + seq_start * delta_batch_stride
|
|
)
|
|
out_base = out_ptr + dim_idx * out_d_stride + seq_start * out_batch_stride
|
|
B_base = B_ptr + group_idx * B_group_stride + seq_start * B_batch_stride
|
|
C_base = C_ptr + group_idx * C_group_stride + seq_start * C_batch_stride
|
|
else:
|
|
u_base = u_ptr + batch_idx * u_batch_stride + dim_idx * u_d_stride
|
|
delta_base = (
|
|
delta_ptr + batch_idx * delta_batch_stride + dim_idx * delta_d_stride
|
|
)
|
|
out_base = out_ptr + batch_idx * out_batch_stride + dim_idx * out_d_stride
|
|
B_base = B_ptr + batch_idx * B_batch_stride + group_idx * B_group_stride
|
|
C_base = C_ptr + batch_idx * C_batch_stride + group_idx * C_group_stride
|
|
|
|
if HAS_Z:
|
|
if IS_VARLEN:
|
|
z_base = z_ptr + dim_idx * z_d_stride + seq_start * z_batch_stride
|
|
out_z_base = (
|
|
out_z_ptr + dim_idx * out_z_d_stride + seq_start * out_z_batch_stride
|
|
)
|
|
else:
|
|
z_base = z_ptr + batch_idx * z_batch_stride + dim_idx * z_d_stride
|
|
out_z_base = (
|
|
out_z_ptr + batch_idx * out_z_batch_stride + dim_idx * out_z_d_stride
|
|
)
|
|
|
|
# Determine chunk boundaries for APC mode
|
|
if CACHE_ENABLED:
|
|
last_chunk_idx = tl.load(last_chunk_indices_ptr + batch_idx).to(tl.int32)
|
|
if batch_idx == 0:
|
|
first_chunk_idx = 0
|
|
else:
|
|
first_chunk_idx = (
|
|
tl.load(last_chunk_indices_ptr + batch_idx - 1).to(tl.int32) + 1
|
|
)
|
|
n_chunks = last_chunk_idx - first_chunk_idx + 1
|
|
first_chunk_tokens = tl.load(cu_chunk_seqlen_ptr + first_chunk_idx + 1).to(
|
|
tl.int32
|
|
) - tl.load(cu_chunk_seqlen_ptr + first_chunk_idx).to(tl.int32)
|
|
block_idx_first = tl.load(block_idx_first_ptr + batch_idx).to(tl.int32)
|
|
chunk_start_offset = 0
|
|
if n_chunks > 1 and first_chunk_tokens < block_size:
|
|
chunk_start_offset = block_size - first_chunk_tokens
|
|
current_position = block_idx_first * block_size + chunk_start_offset
|
|
else:
|
|
n_chunks = 1
|
|
first_chunk_idx = 0
|
|
|
|
# Sequential scan over the sequence
|
|
tokens_processed = 0
|
|
for chunk in range(0, n_chunks if CACHE_ENABLED else 1):
|
|
if CACHE_ENABLED:
|
|
chunk_tokens = tl.load(
|
|
cu_chunk_seqlen_ptr + first_chunk_idx + chunk + 1
|
|
).to(tl.int32) - tl.load(cu_chunk_seqlen_ptr + first_chunk_idx + chunk).to(
|
|
tl.int32
|
|
)
|
|
else:
|
|
chunk_tokens = actual_seqlen
|
|
|
|
for local_pos in range(chunk_tokens):
|
|
pos = tokens_processed + local_pos
|
|
# Load u value
|
|
u_val = tl.load(u_base + pos).to(tl.float32)
|
|
|
|
# Load delta value
|
|
delta_val = tl.load(delta_base + pos).to(tl.float32)
|
|
|
|
# Apply delta bias
|
|
if HAS_DELTA_BIAS:
|
|
delta_val = delta_val + delta_bias_val
|
|
|
|
# Apply softplus
|
|
if delta_softplus:
|
|
delta_val = _softplus(delta_val)
|
|
|
|
delta_u = delta_val * u_val
|
|
|
|
# Compute dA = exp(delta * A) for all dstate elements
|
|
dA = tl.exp(delta_val * A_vals)
|
|
|
|
# Load B values for this position
|
|
B_vals = tl.load(
|
|
B_base + dstate_offs * B_dstate_stride + pos,
|
|
mask=dstate_mask,
|
|
other=0.0,
|
|
).to(tl.float32)
|
|
|
|
# Load C values for this position
|
|
C_vals = tl.load(
|
|
C_base + dstate_offs * C_dstate_stride + pos,
|
|
mask=dstate_mask,
|
|
other=0.0,
|
|
).to(tl.float32)
|
|
|
|
# Update state: state = dA * state + delta * u * B
|
|
state = dA * state + delta_u * B_vals
|
|
|
|
# Compute output: out = sum(state * C) + D * u
|
|
out_val = tl.sum(state * C_vals, axis=0)
|
|
if HAS_D:
|
|
out_val = out_val + D_val * u_val
|
|
|
|
# Store output
|
|
tl.store(out_base + pos, out_val.to(out_ptr.dtype.element_ty))
|
|
|
|
if HAS_Z:
|
|
z_val = tl.load(z_base + pos).to(tl.float32)
|
|
out_z_val = out_val * z_val / (1.0 + tl.exp(-z_val))
|
|
tl.store(
|
|
out_z_base + pos,
|
|
out_z_val.to(out_z_ptr.dtype.element_ty),
|
|
)
|
|
|
|
tokens_processed += chunk_tokens
|
|
|
|
# Store intermediate state for APC mode
|
|
if CACHE_ENABLED:
|
|
if chunk == n_chunks - 1:
|
|
store_slot = tl.load(
|
|
cache_indices_ptr
|
|
+ batch_idx * cache_indices_stride
|
|
+ tl.load(block_idx_last_ptr + batch_idx).to(tl.int32)
|
|
).to(tl.int64)
|
|
else:
|
|
block_idx_done = (current_position + chunk_tokens - 1) // block_size
|
|
store_slot = tl.load(
|
|
cache_indices_ptr
|
|
+ batch_idx * cache_indices_stride
|
|
+ block_idx_done
|
|
).to(tl.int64)
|
|
|
|
tl.store(
|
|
ssm_states_ptr
|
|
+ store_slot * ssm_batch_stride
|
|
+ dim_idx * ssm_dim_stride
|
|
+ dstate_offs * ssm_dstate_stride,
|
|
state.to(ssm_states_ptr.dtype.element_ty),
|
|
mask=dstate_mask,
|
|
)
|
|
current_position += chunk_tokens
|
|
|
|
# Store final state for non-APC mode
|
|
if not CACHE_ENABLED:
|
|
tl.store(
|
|
ssm_states_ptr
|
|
+ load_cache_slot * ssm_batch_stride
|
|
+ dim_idx * ssm_dim_stride
|
|
+ dstate_offs * ssm_dstate_stride,
|
|
state.to(ssm_states_ptr.dtype.element_ty),
|
|
mask=dstate_mask,
|
|
)
|
|
|
|
|
|
# Global flag to ensure ops are registered only once
|
|
_OPS_REGISTERED = False
|
|
|
|
|
|
class xpu_ops:
|
|
@staticmethod
|
|
@torch.compile
|
|
def dynamic_per_token_int8_quant_ref(
|
|
input: torch.Tensor, use_sym_quant: bool, bits: int
|
|
):
|
|
original_sizes = input.size()
|
|
# view is not safe in torch.compile if input is not contiguous
|
|
input = input.reshape(
|
|
-1, original_sizes[-1]
|
|
) # Flatten except for the last dimension
|
|
qmin = -(2 ** (bits - 1)) if use_sym_quant else 0
|
|
qmax = 2 ** (bits - 1) - 1 if use_sym_quant else 2**bits - 1
|
|
min_val = torch.min(input, dim=-1)[0].to(dtype=torch.float32).unsqueeze(-1)
|
|
max_val = torch.max(input, dim=-1)[0].to(dtype=torch.float32).unsqueeze(-1)
|
|
if use_sym_quant:
|
|
scale = (
|
|
torch.maximum(torch.abs(min_val), torch.abs(max_val)) / qmax
|
|
).clamp(min=1e-5)
|
|
zero_point = torch.zeros_like(scale).to(dtype=torch.int32)
|
|
else:
|
|
scale = ((max_val - min_val) / qmax).clamp(min=1e-5)
|
|
zero_point = -1 * torch.round(min_val / scale).to(dtype=torch.int32)
|
|
scale = scale.to(dtype=input.dtype)
|
|
quantized = torch.clamp(
|
|
torch.round(input / scale.to(dtype=torch.float32) + zero_point),
|
|
qmin,
|
|
qmax,
|
|
).to(dtype=torch.int8 if use_sym_quant else torch.uint8)
|
|
return (
|
|
quantized.view(original_sizes),
|
|
scale.view(original_sizes[:-1] + (1,)),
|
|
zero_point.view(original_sizes[:-1] + (1,)),
|
|
)
|
|
|
|
@staticmethod
|
|
def flash_attn_varlen_func(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
cu_seqlens_q: torch.Tensor,
|
|
max_seqlen_q: int,
|
|
max_seqlen_k: int,
|
|
softmax_scale: float | None = None,
|
|
causal: bool = False,
|
|
out: torch.Tensor | None = None,
|
|
block_table: torch.Tensor | None = None,
|
|
alibi_slopes: torch.Tensor | None = None,
|
|
window_size: list[int] | None = None,
|
|
softcap: float | None = 0.0,
|
|
seqused_k: torch.Tensor | None = None,
|
|
cu_seqlens_k: torch.Tensor | None = None,
|
|
# passed in qwen vl
|
|
dropout_p: float = 0.0,
|
|
# The following parameters are not used in xpu kernel currently,
|
|
# we keep API compatible to CUDA's.
|
|
scheduler_metadata=None,
|
|
fa_version: int = 2,
|
|
q_descale=None,
|
|
k_descale=None,
|
|
v_descale=None,
|
|
num_splits=0,
|
|
return_softmax_lse: bool | None = False,
|
|
s_aux: torch.Tensor | None = None,
|
|
return_attn_probs: bool | None = False,
|
|
dynamic_causal: torch.Tensor | None = None,
|
|
mask_mod: Callable | None = None,
|
|
aux_tensors: list | None = None,
|
|
**kwargs,
|
|
):
|
|
assert cu_seqlens_k is not None or seqused_k is not None, (
|
|
"cu_seqlens_k or seqused_k must be provided"
|
|
)
|
|
assert cu_seqlens_k is None or seqused_k is None, (
|
|
"cu_seqlens_k and seqused_k cannot be provided at the same time"
|
|
)
|
|
assert block_table is None or seqused_k is not None, (
|
|
"when enable block_table, seqused_k is needed"
|
|
)
|
|
assert block_table is not None or cu_seqlens_k is not None, (
|
|
"when block_table is disabled, cu_seqlens_k is needed"
|
|
)
|
|
if out is None:
|
|
out = torch.empty(q.shape, dtype=q.dtype, device=q.device)
|
|
real_window_size: tuple[int, int]
|
|
if window_size is None:
|
|
real_window_size = (-1, -1)
|
|
else:
|
|
assert len(window_size) == 2
|
|
real_window_size = (window_size[0], window_size[1]) # noqa: F841
|
|
|
|
return flash_attn_varlen_func(
|
|
out=out,
|
|
q=q,
|
|
k=k,
|
|
v=v,
|
|
cu_seqlens_q=cu_seqlens_q,
|
|
cu_seqlens_k=cu_seqlens_k,
|
|
seqused_k=seqused_k,
|
|
max_seqlen_q=max_seqlen_q,
|
|
max_seqlen_k=max_seqlen_k,
|
|
softmax_scale=softmax_scale,
|
|
causal=causal,
|
|
block_table=block_table,
|
|
s_aux=s_aux,
|
|
window_size=real_window_size,
|
|
# alibi_slopes = alibi_slopes,
|
|
# softcap=softcap,
|
|
return_softmax_lse=return_softmax_lse,
|
|
q_descale=q_descale,
|
|
k_descale=k_descale,
|
|
v_descale=v_descale,
|
|
)
|
|
|
|
@staticmethod
|
|
def get_scheduler_metadata(
|
|
batch_size,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
num_heads_q,
|
|
num_heads_kv,
|
|
headdim,
|
|
cache_seqlens: torch.Tensor,
|
|
qkv_dtype=torch.bfloat16,
|
|
headdim_v=None,
|
|
cu_seqlens_q: torch.Tensor | None = None,
|
|
cu_seqlens_k_new: torch.Tensor | None = None,
|
|
cache_leftpad: torch.Tensor | None = None,
|
|
page_size: int | None = None,
|
|
max_seqlen_k_new=0,
|
|
causal=False,
|
|
window_size=(-1, -1), # -1 means infinite context window
|
|
has_softcap=False,
|
|
num_splits=0, # Can be tuned for speed
|
|
pack_gqa=None, # Can be tuned for speed
|
|
sm_margin=0, # Can be tuned if some SMs are used for communication
|
|
) -> None:
|
|
logger.warning_once(
|
|
"get_scheduler_metadata is not implemented for xpu_ops, returning None."
|
|
)
|
|
return None
|
|
|
|
@staticmethod
|
|
def selective_scan_fwd(
|
|
u: torch.Tensor,
|
|
delta: torch.Tensor,
|
|
A: torch.Tensor,
|
|
B: torch.Tensor,
|
|
C: torch.Tensor,
|
|
D_: torch.Tensor | None,
|
|
z_: torch.Tensor | None,
|
|
delta_bias_: torch.Tensor | None,
|
|
delta_softplus: bool,
|
|
query_start_loc: torch.Tensor | None,
|
|
cache_indices: torch.Tensor | None,
|
|
has_initial_state: torch.Tensor | None,
|
|
ssm_states: torch.Tensor,
|
|
null_block_id: int,
|
|
block_size: int = 1024,
|
|
block_idx_first_scheduled_token: torch.Tensor | None = None,
|
|
block_idx_last_scheduled_token: torch.Tensor | None = None,
|
|
initial_state_idx: torch.Tensor | None = None,
|
|
cu_chunk_seqlen: torch.Tensor | None = None,
|
|
last_chunk_indices: torch.Tensor | None = None,
|
|
) -> None:
|
|
varlen = query_start_loc is not None
|
|
batch_size = (
|
|
(query_start_loc.shape[0] - 1)
|
|
if query_start_loc is not None
|
|
else u.shape[0]
|
|
)
|
|
dim = u.shape[0] if varlen else u.shape[1]
|
|
total_seqlen = u.shape[1] if varlen else u.shape[2]
|
|
dstate = A.size(1)
|
|
n_groups = B.size(0) if varlen else B.size(1)
|
|
dim_ngroups_ratio = dim // n_groups
|
|
|
|
has_z = z_ is not None
|
|
has_D = D_ is not None
|
|
has_delta_bias = delta_bias_ is not None
|
|
has_cache_indices = cache_indices is not None
|
|
cache_enabled = block_idx_first_scheduled_token is not None
|
|
|
|
# out and out_z alias delta and z respectively
|
|
out = delta
|
|
out_z = z_ if z_ is not None else delta # won't be used if not has_z
|
|
|
|
BLOCK_DSTATE = triton.next_power_of_2(dstate)
|
|
|
|
# Compute strides
|
|
if varlen:
|
|
u_batch_stride = u.stride(1)
|
|
u_d_stride = u.stride(0)
|
|
delta_batch_stride = delta.stride(1)
|
|
delta_d_stride = delta.stride(0)
|
|
B_batch_stride = B.stride(2)
|
|
B_group_stride = B.stride(0)
|
|
B_dstate_stride = B.stride(1)
|
|
C_batch_stride = C.stride(2)
|
|
C_group_stride = C.stride(0)
|
|
C_dstate_stride = C.stride(1)
|
|
out_batch_stride = out.stride(1)
|
|
out_d_stride = out.stride(0)
|
|
if z_ is not None:
|
|
z_batch_stride = z_.stride(1)
|
|
z_d_stride = z_.stride(0)
|
|
out_z_batch_stride = out_z.stride(1)
|
|
out_z_d_stride = out_z.stride(0)
|
|
else:
|
|
z_batch_stride = 0
|
|
z_d_stride = 0
|
|
out_z_batch_stride = 0
|
|
out_z_d_stride = 0
|
|
else:
|
|
u_batch_stride = u.stride(0)
|
|
u_d_stride = u.stride(1)
|
|
delta_batch_stride = delta.stride(0)
|
|
delta_d_stride = delta.stride(1)
|
|
B_batch_stride = B.stride(0)
|
|
B_group_stride = B.stride(1)
|
|
B_dstate_stride = B.stride(2)
|
|
C_batch_stride = C.stride(0)
|
|
C_group_stride = C.stride(1)
|
|
C_dstate_stride = C.stride(2)
|
|
out_batch_stride = out.stride(0)
|
|
out_d_stride = out.stride(1)
|
|
if z_ is not None:
|
|
z_batch_stride = z_.stride(0)
|
|
z_d_stride = z_.stride(1)
|
|
out_z_batch_stride = out_z.stride(0)
|
|
out_z_d_stride = out_z.stride(1)
|
|
else:
|
|
z_batch_stride = 0
|
|
z_d_stride = 0
|
|
out_z_batch_stride = 0
|
|
out_z_d_stride = 0
|
|
|
|
ssm_batch_stride = ssm_states.stride(0)
|
|
ssm_dim_stride = ssm_states.stride(1)
|
|
ssm_dstate_stride = ssm_states.stride(2)
|
|
cache_indices_stride = (
|
|
cache_indices.stride(0) if cache_indices is not None else 0
|
|
)
|
|
|
|
grid = (batch_size, dim)
|
|
_selective_scan_fwd_kernel[grid](
|
|
u,
|
|
delta,
|
|
A,
|
|
B,
|
|
C,
|
|
D_ if has_D else u, # dummy, won't be dereferenced
|
|
z_ if has_z else u, # dummy
|
|
delta_bias_ if has_delta_bias else u, # dummy
|
|
out,
|
|
out_z,
|
|
ssm_states,
|
|
query_start_loc if varlen else u, # dummy
|
|
cache_indices if has_cache_indices else u, # dummy
|
|
has_initial_state,
|
|
# APC pointers
|
|
block_idx_first_scheduled_token if cache_enabled else u,
|
|
block_idx_last_scheduled_token if cache_enabled else u,
|
|
initial_state_idx if cache_enabled else u,
|
|
cu_chunk_seqlen if cache_enabled else u,
|
|
last_chunk_indices if cache_enabled else u,
|
|
# Dimensions
|
|
batch_size,
|
|
dim,
|
|
total_seqlen,
|
|
dstate,
|
|
n_groups,
|
|
dim_ngroups_ratio,
|
|
# Strides
|
|
u_batch_stride,
|
|
u_d_stride,
|
|
delta_batch_stride,
|
|
delta_d_stride,
|
|
A.stride(0),
|
|
A.stride(1),
|
|
B_batch_stride,
|
|
B_group_stride,
|
|
B_dstate_stride,
|
|
C_batch_stride,
|
|
C_group_stride,
|
|
C_dstate_stride,
|
|
z_batch_stride,
|
|
z_d_stride,
|
|
out_batch_stride,
|
|
out_d_stride,
|
|
out_z_batch_stride,
|
|
out_z_d_stride,
|
|
ssm_batch_stride,
|
|
ssm_dim_stride,
|
|
ssm_dstate_stride,
|
|
cache_indices_stride,
|
|
null_block_id,
|
|
block_size,
|
|
# Compile-time constants
|
|
delta_softplus=delta_softplus,
|
|
HAS_D=has_D,
|
|
HAS_Z=has_z,
|
|
HAS_DELTA_BIAS=has_delta_bias,
|
|
IS_VARLEN=varlen,
|
|
HAS_CACHE_INDICES=has_cache_indices,
|
|
CACHE_ENABLED=cache_enabled,
|
|
BLOCK_DSTATE=BLOCK_DSTATE,
|
|
)
|
|
|
|
@staticmethod
|
|
def register_ops_once() -> None:
|
|
global _OPS_REGISTERED
|
|
if not _OPS_REGISTERED:
|
|
# register all the custom ops here
|
|
direct_register_custom_op(
|
|
op_name="xpu_ops_deepseek_scaling_rope",
|
|
op_func=_xpu_ops_deepseek_scaling_rope_impl,
|
|
mutates_args=[],
|
|
fake_impl=_xpu_ops_deepseek_scaling_rope_fake,
|
|
dispatch_key=current_platform.dispatch_key,
|
|
)
|
|
|
|
direct_register_custom_op(
|
|
op_name="xpu_mxfp8_quantize",
|
|
op_func=_xpu_mxfp8_quantize_impl,
|
|
fake_impl=_xpu_mxfp8_quantize_fake,
|
|
)
|
|
|
|
direct_register_custom_op(
|
|
op_name="xpu_mxfp4_quantize",
|
|
op_func=_xpu_mxfp4_quantize_impl,
|
|
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,
|
|
mutates_args=["core_attn_out", "z"],
|
|
fake_impl=_gdn_attention_core_xpu_fake,
|
|
)
|
|
|
|
direct_register_custom_op(
|
|
op_name="xpu_topk_topp_sampler",
|
|
op_func=_topk_topp_sample_impl,
|
|
fake_impl=_topk_topp_sample_fake,
|
|
)
|
|
|
|
_OPS_REGISTERED = True
|
|
|
|
|
|
xpu_ops.register_ops_once()
|