# 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()