| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Callable |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | import torch |
| 8 | from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func |
| 9 | from vllm_xpu_kernels.rotary import apply_rotary_emb |
| 10 | |
| 11 | from vllm.compilation.breakable_cudagraph import eager_break_during_capture |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.platforms import current_platform |
| 14 | from vllm.triton_utils import tl, triton |
| 15 | from vllm.utils.torch_utils import direct_register_custom_op |
| 16 | |
| 17 | logger = init_logger(__name__) |
| 18 | |
| 19 | if TYPE_CHECKING: |
| 20 | |
| 21 | def register_fake(fn): |
| 22 | return lambda name: fn |
| 23 | else: |
| 24 | try: |
| 25 | from torch.library import register_fake |
| 26 | except ImportError: |
| 27 | from torch.library import impl_abstract as register_fake |
| 28 | |
| 29 | if hasattr(torch.ops._xpu_C, "fp8_gemm"): |
| 30 | |
| 31 | @register_fake("_xpu_C::fp8_gemm") |
| 32 | def _fp8_gemm_fake( |
| 33 | q_input: torch.Tensor, |
| 34 | q_weight: torch.Tensor, |
| 35 | out_dtype: torch.dtype, |
| 36 | input_scales: torch.Tensor, |
| 37 | weight_scale: torch.Tensor, |
| 38 | bias: torch.Tensor | None = None, |
| 39 | ) -> torch.Tensor: |
| 40 | input_2d = q_input.view(-1, q_input.shape[-1]) |
| 41 | M = input_2d.size(0) |
| 42 | N = q_weight.size(1) |
| 43 | return torch.empty((M, N), dtype=out_dtype, device=q_input.device) |
| 44 | |
| 45 | |
| 46 | if hasattr(torch.ops._xpu_C, "fp8_gemm_w8a16"): |
| 47 | |
| 48 | @register_fake("_xpu_C::fp8_gemm_w8a16") |
| 49 | def _fp8_gemm_w8a16_fake( |
| 50 | input: torch.Tensor, |
| 51 | q_weight: torch.Tensor, |
| 52 | weight_scale: torch.Tensor, |
| 53 | bias: torch.Tensor | None = None, |
| 54 | ) -> torch.Tensor: |
| 55 | input_2d = input.view(-1, input.shape[-1]) |
| 56 | M = input_2d.size(0) |
| 57 | N = q_weight.size(1) |
| 58 | return torch.empty((M, N), dtype=input.dtype, device=input.device) |
| 59 | |
| 60 | |
| 61 | if hasattr(torch.ops._xpu_C, "fp4_gemm"): |
| 62 | |
| 63 | @register_fake("_xpu_C::fp4_gemm") |
| 64 | def _fp4_gemm_fake( |
| 65 | input: torch.Tensor, |
| 66 | weight: torch.Tensor, |
| 67 | scale_act: torch.Tensor, |
| 68 | scale_wei: torch.Tensor, |
| 69 | out_dtype: torch.dtype | None = None, |
| 70 | bias: torch.Tensor | None = None, |
| 71 | ) -> torch.Tensor: |
| 72 | input_2d = input.view(-1, input.shape[-1]) |
| 73 | M = input_2d.size(0) |
| 74 | N = weight.size(-1) |
| 75 | result_dtype = out_dtype if out_dtype is not None else torch.get_default_dtype() |
| 76 | return torch.empty((M, N), dtype=result_dtype, device=input.device) |
| 77 | |
| 78 | |
| 79 | if hasattr(torch.ops._xpu_C, "int4_gemm_w4a8"): |
| 80 | |
| 81 | @register_fake("_xpu_C::int4_gemm_w4a8") |
| 82 | def _int4_gemm_w4a8_fake( |
| 83 | input: torch.Tensor, |
| 84 | input_scales: torch.Tensor, |
| 85 | input_zero_points: torch.Tensor, |
| 86 | q_weight: torch.Tensor, |
| 87 | weight_scale: torch.Tensor, |
| 88 | weight_zp: torch.Tensor, |
| 89 | group_size: int, |
| 90 | g_idx: torch.Tensor | None = None, |
| 91 | bias: torch.Tensor | None = None, |
| 92 | ) -> torch.Tensor: |
| 93 | input_2d = input.view(-1, input.shape[-1]) |
| 94 | M = input_2d.size(0) |
| 95 | N = q_weight.size(1) |
| 96 | return torch.empty((M, N), dtype=torch.float16, device=input.device) |
| 97 | |
| 98 | |
| 99 | if hasattr(torch.ops._xpu_C, "int4_gemm_w4a16"): |
| 100 | |
| 101 | @register_fake("_xpu_C::int4_gemm_w4a16") |
| 102 | def _int4_gemm_w4a16_fake( |
| 103 | input: torch.Tensor, |
| 104 | q_weight: torch.Tensor, |
| 105 | bias: torch.Tensor | None, |
| 106 | weight_scale: torch.Tensor, |
| 107 | qzeros: torch.Tensor, |
| 108 | group_size: int, |
| 109 | group_idx: torch.Tensor | None = None, |
| 110 | ) -> torch.Tensor: |
| 111 | input_2d = input.view(-1, input.shape[-1]) |
| 112 | M = input_2d.size(0) |
| 113 | N = q_weight.size(1) |
| 114 | return torch.empty((M, N), dtype=input.dtype, device=input.device) |
| 115 | |
| 116 | |
| 117 | def _gemma_rms_norm_impl( |
| 118 | out: torch.Tensor, |
| 119 | input: torch.Tensor, |
| 120 | weight: torch.Tensor, |
| 121 | epsilon: float, |
| 122 | ) -> None: |
| 123 | # GemmaRMSNorm: computes out = (x_normed_fp32 * (1 + weight.float()) |
| 124 | # ).to(dtype) with a raw (bf16/fp16) weight; the +1 offset and fp32 |
| 125 | # multiply are done in-kernel. See vllm-xpu-kernels gemma_rms_norm. |
| 126 | torch.ops._C.gemma_rms_norm(out, input, weight, epsilon) |
| 127 | |
| 128 | |
| 129 | def _fused_add_gemma_rms_norm_impl( |
| 130 | input: torch.Tensor, |
| 131 | residual: torch.Tensor, |
| 132 | weight: torch.Tensor, |
| 133 | epsilon: float, |
| 134 | ) -> None: |
| 135 | torch.ops._C.fused_add_gemma_rms_norm(input, residual, weight, epsilon) |
| 136 | |
| 137 | |
| 138 | def _layer_norm_impl( |
| 139 | out: torch.Tensor, |
| 140 | input: torch.Tensor, |
| 141 | weight: torch.Tensor, |
| 142 | bias: torch.Tensor, |
| 143 | epsilon: float, |
| 144 | ) -> None: |
| 145 | torch.ops._C.layer_norm(out, input, weight, bias, epsilon) |
| 146 | |
| 147 | |
| 148 | def _fused_add_layer_norm_impl( |
| 149 | input: torch.Tensor, |
| 150 | residual: torch.Tensor, |
| 151 | weight: torch.Tensor, |
| 152 | bias: torch.Tensor, |
| 153 | epsilon: float, |
| 154 | ) -> None: |
| 155 | torch.ops._C.fused_add_layer_norm(input, residual, weight, bias, epsilon) |
| 156 | |
| 157 | |
| 158 | def _nemotron_layer_norm_impl( |
| 159 | out: torch.Tensor, |
| 160 | input: torch.Tensor, |
| 161 | weight: torch.Tensor, |
| 162 | bias: torch.Tensor, |
| 163 | epsilon: float, |
| 164 | ) -> None: |
| 165 | torch.ops._C.nemotron_layer_norm(out, input, weight, bias, epsilon) |
| 166 | |
| 167 | |
| 168 | def _fused_add_nemotron_layer_norm_impl( |
| 169 | input: torch.Tensor, |
| 170 | residual: torch.Tensor, |
| 171 | weight: torch.Tensor, |
| 172 | bias: torch.Tensor, |
| 173 | epsilon: float, |
| 174 | ) -> None: |
| 175 | torch.ops._C.fused_add_nemotron_layer_norm(input, residual, weight, bias, epsilon) |
| 176 | |
| 177 | |
| 178 | def _gdn_attention_core_xpu_impl( |
| 179 | core_attn_out: torch.Tensor, |
| 180 | z: torch.Tensor, |
| 181 | projected_states_qkvz: torch.Tensor, |
| 182 | projected_states_ba: torch.Tensor, |
| 183 | layer_name: str, |
| 184 | ) -> None: |
| 185 | """Custom op wrapping the XPU SYCL GDN kernel for torch.compile.""" |
| 186 | from vllm.forward_context import get_forward_context |
| 187 | from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata |
| 188 | |
| 189 | forward_context = get_forward_context() |
| 190 | self = forward_context.no_compile_layers[layer_name] |
| 191 | attn_metadata_raw = forward_context.attn_metadata |
| 192 | |
| 193 | if attn_metadata_raw is None: |
| 194 | return |
| 195 | |
| 196 | assert isinstance(attn_metadata_raw, dict) |
| 197 | attn_metadata = attn_metadata_raw[self.prefix] |
| 198 | assert isinstance(attn_metadata, GDNAttentionMetadata) |
| 199 | |
| 200 | num_actual_tokens = attn_metadata.num_actual_tokens |
| 201 | num_accepted_tokens = attn_metadata.num_accepted_tokens |
| 202 | |
| 203 | num_prefills = attn_metadata.num_prefills |
| 204 | num_decodes = attn_metadata.num_decodes |
| 205 | num_spec_decodes = attn_metadata.num_spec_decodes |
| 206 | |
| 207 | has_initial_state = attn_metadata.has_initial_state |
| 208 | |
| 209 | non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc |
| 210 | non_spec_token_indx = attn_metadata.non_spec_token_indx |
| 211 | non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501 |
| 212 | non_spec_state_indices_tensor = ( |
| 213 | non_spec_state_indices_tensor.contiguous() |
| 214 | if non_spec_state_indices_tensor is not None |
| 215 | else None |
| 216 | ) |
| 217 | |
| 218 | spec_query_start_loc = attn_metadata.spec_query_start_loc |
| 219 | spec_token_indx = attn_metadata.spec_token_indx |
| 220 | spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501 |
| 221 | |
| 222 | spec_sequence_masks = attn_metadata.spec_sequence_masks |
| 223 | if spec_sequence_masks is not None: |
| 224 | if non_spec_token_indx is not None: |
| 225 | non_spec_token_indx = non_spec_token_indx.to(torch.int32) |
| 226 | if spec_token_indx is not None: |
| 227 | spec_token_indx = spec_token_indx.to(torch.int32) |
| 228 | |
| 229 | conv_weights = self.conv1d.weight.view( |
| 230 | self.conv1d.weight.size(0), self.conv1d.weight.size(2) |
| 231 | ) |
| 232 | |
| 233 | torch.ops._xpu_C.gdn_attention( |
| 234 | core_attn_out, |
| 235 | z, |
| 236 | projected_states_qkvz, |
| 237 | projected_states_ba, |
| 238 | self.num_k_heads, |
| 239 | self.num_v_heads, |
| 240 | self.head_k_dim, |
| 241 | self.head_v_dim, |
| 242 | conv_state=self.kv_cache[0], |
| 243 | ssm_state=self.kv_cache[1], |
| 244 | conv_weights=conv_weights, |
| 245 | conv_bias=self.conv1d.bias, |
| 246 | activation=self.activation, |
| 247 | A_log=self.A_log, |
| 248 | dt_bias=self.dt_bias, |
| 249 | num_prefills=num_prefills, # type: ignore[attr-defined] |
| 250 | num_decodes=num_decodes, # type: ignore[attr-defined] |
| 251 | num_spec_decodes=num_spec_decodes, # type: ignore[attr-defined] |
| 252 | has_initial_state=has_initial_state, # type: ignore[attr-defined] |
| 253 | non_spec_query_start_loc=non_spec_query_start_loc, # type: ignore[attr-defined] |
| 254 | non_spec_token_indx=non_spec_token_indx, # type: ignore[attr-defined] |
| 255 | non_spec_state_indices_tensor=non_spec_state_indices_tensor, # type: ignore[attr-defined] |
| 256 | spec_query_start_loc=spec_query_start_loc, # type: ignore[attr-defined] |
| 257 | spec_token_indx=spec_token_indx, # type: ignore[attr-defined] |
| 258 | spec_state_indices_tensor=spec_state_indices_tensor, |
| 259 | num_accepted_tokens=num_accepted_tokens, # type: ignore[attr-defined] |
| 260 | num_actual_tokens=num_actual_tokens, # type: ignore[attr-defined] |
| 261 | tp_size=self.tp_size, |
| 262 | reorder_input=not self.gqa_interleaved_layout, |
| 263 | ) |
| 264 | |
| 265 | |
| 266 | def _xpu_ops_deepseek_scaling_rope_impl( |
| 267 | positions: torch.Tensor, |
| 268 | query: torch.Tensor, |
| 269 | key: torch.Tensor | None, |
| 270 | offsets: torch.Tensor | None, |
| 271 | cos_sin_cache: torch.Tensor | None, |
| 272 | rotary_dim: int, |
| 273 | is_neox_style: bool, |
| 274 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 275 | assert key is not None |
| 276 | return torch.ops._xpu_C.deepseek_scaling_rope( |
| 277 | positions, query, key, offsets, cos_sin_cache, rotary_dim, is_neox_style |
| 278 | ) |
| 279 | |
| 280 | |
| 281 | def _xpu_ops_deepseek_scaling_rope_fake( |
| 282 | positions: torch.Tensor, |
| 283 | query: torch.Tensor, |
| 284 | key: torch.Tensor | None, |
| 285 | offsets: torch.Tensor | None, |
| 286 | cos_sin_cache: torch.Tensor | None, |
| 287 | rotary_dim: int, |
| 288 | is_neox_style: bool, |
| 289 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 290 | return query, key |
| 291 | |
| 292 | |
| 293 | def _xpu_apply_rotary_emb_impl( |
| 294 | x: torch.Tensor, |
| 295 | cos: torch.Tensor, |
| 296 | sin: torch.Tensor, |
| 297 | is_neox_style: bool, |
| 298 | ) -> torch.Tensor: |
| 299 | return apply_rotary_emb(x, cos, sin, is_neox_style) |
| 300 | |
| 301 | |
| 302 | def _xpu_apply_rotary_emb_fake( |
| 303 | x: torch.Tensor, |
| 304 | cos: torch.Tensor, |
| 305 | sin: torch.Tensor, |
| 306 | is_neox_style: bool, |
| 307 | ) -> torch.Tensor: |
| 308 | return torch.empty_like(x) |
| 309 | |
| 310 | |
| 311 | def _xpu_fp8_bmm_impl( |
| 312 | a: torch.Tensor, |
| 313 | b: torch.Tensor, |
| 314 | out_dtype: torch.dtype, |
| 315 | a_scale: torch.Tensor, |
| 316 | b_scale: torch.Tensor, |
| 317 | bias: torch.Tensor | None, |
| 318 | ) -> torch.Tensor: |
| 319 | """XPU FP8 batched GEMM implementation for ``torch.ops.vllm.xpu_fp8_bmm``. |
| 320 | |
| 321 | Computes batched matrix multiplication over the leading group dimension: |
| 322 | ``[G, M, K] @ [G, K, N] -> [G, M, N]``. |
| 323 | |
| 324 | Args: |
| 325 | a: FP8 activation tensor with shape ``[G, M, K]``. |
| 326 | Does not need to be contiguous. |
| 327 | b: FP8 weight tensor with shape ``[G, K, N]``. |
| 328 | Does not need to be contiguous. |
| 329 | out_dtype: Output dtype accepted by the kernel (typically |
| 330 | ``torch.bfloat16`` for the DeepSeek-V4 O-proj path). |
| 331 | a_scale: Activation scale tensor for ``a``. |
| 332 | In current DeepSeek-V4 XPU usage it is block-scaled with shape |
| 333 | ``[G, M, K/bs]`` (``bs`` is the quant block size, e.g. 128). |
| 334 | Must be contiguous. |
| 335 | b_scale: Weight scale tensor for ``b``. |
| 336 | In current DeepSeek-V4 XPU usage it is block-scaled with shape |
| 337 | ``[G, K/bs, N/bs]`` (``bs`` is the quant block size, e.g. 128). |
| 338 | Must be contiguous. |
| 339 | bias: Optional bias tensor. Pass ``None`` when no bias is required. |
| 340 | |
| 341 | Returns: |
| 342 | Output tensor with shape ``[G, M, N]`` and dtype ``out_dtype``. |
| 343 | |
| 344 | Notes: |
| 345 | This implementation centralizes access to |
| 346 | ``torch.ops._xpu_C.fp8_bmm``. Both scales must be contiguous, while |
| 347 | ``a`` and ``b`` may be non-contiguous views. |
| 348 | |
| 349 | """ |
| 350 | return torch.ops._xpu_C.fp8_bmm(a, b, out_dtype, a_scale, b_scale, bias) |
| 351 | |
| 352 | |
| 353 | def _xpu_fp8_bmm_fake( |
| 354 | a: torch.Tensor, |
| 355 | b: torch.Tensor, |
| 356 | out_dtype: torch.dtype, |
| 357 | a_scale: torch.Tensor, |
| 358 | b_scale: torch.Tensor, |
| 359 | bias: torch.Tensor | None, |
| 360 | ) -> torch.Tensor: |
| 361 | # [G, M, K] @ [G, K, N] => [G, M, N] |
| 362 | return torch.empty( |
| 363 | (a.shape[0], a.shape[1], b.shape[2]), |
| 364 | dtype=out_dtype, |
| 365 | device=a.device, |
| 366 | ) |
| 367 | |
| 368 | |
| 369 | def _xpu_fp8_mqa_logits_impl( |
| 370 | q: torch.Tensor, |
| 371 | k_quant: torch.Tensor, |
| 372 | k_scale: torch.Tensor, |
| 373 | weights: torch.Tensor, |
| 374 | cu_seqlen_ks: torch.Tensor, |
| 375 | cu_seqlen_ke: torch.Tensor, |
| 376 | ) -> torch.Tensor: |
| 377 | return torch.ops._xpu_C.fp8_mqa_logits( |
| 378 | q, |
| 379 | k_quant, |
| 380 | k_scale, |
| 381 | weights, |
| 382 | cu_seqlen_ks, |
| 383 | cu_seqlen_ke, |
| 384 | ) |
| 385 | |
| 386 | |
| 387 | def _xpu_fp8_mqa_logits_fake( |
| 388 | q: torch.Tensor, |
| 389 | k_quant: torch.Tensor, |
| 390 | k_scale: torch.Tensor, |
| 391 | weights: torch.Tensor, |
| 392 | cu_seqlen_ks: torch.Tensor, |
| 393 | cu_seqlen_ke: torch.Tensor, |
| 394 | ) -> torch.Tensor: |
| 395 | return torch.empty( |
| 396 | (q.shape[0], k_quant.shape[0]), |
| 397 | dtype=torch.float32, |
| 398 | device=q.device, |
| 399 | ) |
| 400 | |
| 401 | |
| 402 | def _xpu_fp8_paged_mqa_logits_impl( |
| 403 | q: torch.Tensor, |
| 404 | kv_cache: torch.Tensor, |
| 405 | weights: torch.Tensor, |
| 406 | context_lens: torch.Tensor, |
| 407 | block_tables: torch.Tensor, |
| 408 | schedule_metadata: torch.Tensor, |
| 409 | max_model_len: int, |
| 410 | ) -> torch.Tensor: |
| 411 | return torch.ops._xpu_C.fp8_paged_mqa_logits( |
| 412 | q, |
| 413 | kv_cache, |
| 414 | weights, |
| 415 | context_lens, |
| 416 | block_tables, |
| 417 | schedule_metadata, |
| 418 | max_model_len, |
| 419 | ) |
| 420 | |
| 421 | |
| 422 | def _xpu_fp8_paged_mqa_logits_fake( |
| 423 | q: torch.Tensor, |
| 424 | kv_cache: torch.Tensor, |
| 425 | weights: torch.Tensor, |
| 426 | context_lens: torch.Tensor, |
| 427 | block_tables: torch.Tensor, |
| 428 | schedule_metadata: torch.Tensor, |
| 429 | max_model_len: int, |
| 430 | ) -> torch.Tensor: |
| 431 | batch_size, next_n = q.shape[:2] |
| 432 | return torch.empty( |
| 433 | (batch_size * next_n, max_model_len), |
| 434 | dtype=torch.float32, |
| 435 | device=q.device, |
| 436 | ) |
| 437 | |
| 438 | |
| 439 | def _topk_topp_sample_impl( |
| 440 | random_sampled: torch.Tensor, |
| 441 | logits_to_return: torch.Tensor | None, |
| 442 | logits: torch.Tensor, |
| 443 | k: torch.Tensor | None, |
| 444 | p: torch.Tensor | None, |
| 445 | logprobs_mode: str, |
| 446 | seeds: torch.Tensor | None, |
| 447 | lambda_: float = 1.0, |
| 448 | ) -> None: |
| 449 | torch.ops._xpu_C.topk_topp_sampler( |
| 450 | random_sampled, logits_to_return, logits, k, p, logprobs_mode, seeds, lambda_ |
| 451 | ) |
| 452 | return |
| 453 | |
| 454 | |
| 455 | def _topk_topp_sample_fake( |
| 456 | random_sampled: torch.Tensor, |
| 457 | logits_to_return: torch.Tensor | None, |
| 458 | logits: torch.Tensor, |
| 459 | k: torch.Tensor | None, |
| 460 | p: torch.Tensor | None, |
| 461 | logprobs_mode: str, |
| 462 | seeds: torch.Tensor | None, |
| 463 | lambda_: float = 1.0, |
| 464 | ) -> None: |
| 465 | return |
| 466 | |
| 467 | |
| 468 | def _xpu_deepseek_fused_indexer_q_rope_fp8_impl( |
| 469 | index_q: torch.Tensor, |
| 470 | positions: torch.Tensor, |
| 471 | index_q_cos_sin_cache: torch.Tensor, |
| 472 | index_weights: torch.Tensor, |
| 473 | index_weights_softmax_scale: float, |
| 474 | index_weights_head_scale: float, |
| 475 | index_q_fp8: torch.Tensor, |
| 476 | index_weights_out: torch.Tensor, |
| 477 | ) -> None: |
| 478 | """Fused RoPE + FP8 quant of the DeepSeek-V4 sparse-indexer Q (XPU). |
| 479 | Writes ``index_q_fp8`` and ``index_weights_out`` in place (no return). |
| 480 | Requires head_dim=128, rope_dim=64, and num_heads divisible by 2. |
| 481 | |
| 482 | Args: |
| 483 | index_q: (T, H, 128) bfloat16 Q before RoPE. Contiguous. |
| 484 | positions: (T,) int64 absolute token positions. |
| 485 | index_q_cos_sin_cache: (max_pos, 64) float32 RoPE cos/sin cache. |
| 486 | index_weights: (T, H) bfloat16 raw indexer weights. |
| 487 | index_weights_softmax_scale: scalar softmax scale. |
| 488 | index_weights_head_scale: scalar per-head scale. |
| 489 | index_q_fp8: (T, H, 128) fp8 (e4m3) output, preallocated. [written] |
| 490 | index_weights_out: (T, H) float32 output, preallocated; |
| 491 | = index_weights * q_scale * softmax_scale * head_scale (the |
| 492 | per-(token, head) q_scale is folded in here). [written] |
| 493 | |
| 494 | """ |
| 495 | torch.ops._xpu_C.deepseek_fused_indexer_q_rope_fp8( |
| 496 | index_q, |
| 497 | positions, |
| 498 | index_q_cos_sin_cache, |
| 499 | index_weights, |
| 500 | index_weights_softmax_scale, |
| 501 | index_weights_head_scale, |
| 502 | index_q_fp8, |
| 503 | index_weights_out, |
| 504 | ) |
| 505 | |
| 506 | |
| 507 | def _xpu_deepseek_fused_indexer_q_rope_mxfp4_impl( |
| 508 | index_q: torch.Tensor, |
| 509 | positions: torch.Tensor, |
| 510 | index_q_cos_sin_cache: torch.Tensor, |
| 511 | index_weights: torch.Tensor, |
| 512 | index_weights_softmax_scale: float, |
| 513 | index_weights_head_scale: float, |
| 514 | index_q_packed: torch.Tensor, |
| 515 | index_q_scale: torch.Tensor, |
| 516 | index_weights_out: torch.Tensor, |
| 517 | ) -> None: |
| 518 | """Fused RoPE + MXFP4 quant of the DeepSeek-V4 sparse-indexer Q (XPU). |
| 519 | Writes ``index_q_packed``, ``index_q_scale`` and ``index_weights_out`` in |
| 520 | place (no return). Requires head_dim=128, rope_dim=64, and num_heads |
| 521 | divisible by 4. |
| 522 | |
| 523 | Args: |
| 524 | index_q: (T, H, 128) bfloat16 Q before RoPE. Contiguous. |
| 525 | positions: (T,) int64 absolute token positions. |
| 526 | index_q_cos_sin_cache: (max_pos, 64) float32 RoPE cos/sin cache. |
| 527 | index_weights: (T, H) bfloat16 raw indexer weights. |
| 528 | index_weights_softmax_scale: scalar softmax scale. |
| 529 | index_weights_head_scale: scalar per-head scale. |
| 530 | index_q_packed: (T, H, 64) uint8 packed E2M1 nibbles (2 per byte), |
| 531 | preallocated. [written] |
| 532 | index_q_scale: (T, H, 4) uint8 ue8m0 per-block (32-elem) scales, |
| 533 | preallocated. [written] |
| 534 | index_weights_out: (T, H) float32 output, preallocated; |
| 535 | = index_weights * softmax_scale * head_scale (no q_scale folded; |
| 536 | per-block scales live in index_q_scale). [written] |
| 537 | |
| 538 | """ |
| 539 | torch.ops._xpu_C.deepseek_fused_indexer_q_rope_mxfp4( |
| 540 | index_q, |
| 541 | positions, |
| 542 | index_q_cos_sin_cache, |
| 543 | index_weights, |
| 544 | index_weights_softmax_scale, |
| 545 | index_weights_head_scale, |
| 546 | index_q_packed, |
| 547 | index_q_scale, |
| 548 | index_weights_out, |
| 549 | ) |
| 550 | |
| 551 | |
| 552 | def _xpu_mxfp8_quantize_impl( |
| 553 | x: torch.Tensor, dtype: torch.dtype | None = None |
| 554 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 555 | MXFP8_BLOCK_SIZE = 32 |
| 556 | assert x.shape[-1] % MXFP8_BLOCK_SIZE == 0 |
| 557 | if dtype is not None: |
| 558 | assert dtype in (torch.float8_e4m3fn, torch.float8_e5m2), ( |
| 559 | f"Unsupported dtype for xpu_mxfp8_quantize: {dtype}. " |
| 560 | f"Expected torch.float8_e4m3fn or torch.float8_e5m2." |
| 561 | ) |
| 562 | else: |
| 563 | dtype = current_platform.fp8_dtype() |
| 564 | |
| 565 | finfo = torch.finfo(dtype) |
| 566 | fp8_min = finfo.min |
| 567 | fp8_max = finfo.max |
| 568 | eps = 1e-10 |
| 569 | |
| 570 | x_q = torch.empty_like(x, device=x.device, dtype=dtype) |
| 571 | shape = x.shape[:-1] + (x.shape[-1] // MXFP8_BLOCK_SIZE,) |
| 572 | x_s = torch.empty(shape, device=x.device, dtype=torch.float32) |
| 573 | torch.ops._C.per_token_group_fp8_quant( |
| 574 | x, |
| 575 | x_q, |
| 576 | x_s, |
| 577 | MXFP8_BLOCK_SIZE, |
| 578 | eps, |
| 579 | fp8_min, |
| 580 | fp8_max, |
| 581 | True, |
| 582 | False, |
| 583 | False, # dummy_is_scale_transposed, dummy_is_tma_aligned |
| 584 | ) |
| 585 | x_s = x_s.to(torch.float8_e8m0fnu) |
| 586 | return x_q, x_s |
| 587 | |
| 588 | |
| 589 | def _xpu_mxfp8_quantize_fake( |
| 590 | x: torch.Tensor, dtype: torch.dtype | None = None |
| 591 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 592 | if dtype is None: |
| 593 | dtype = current_platform.fp8_dtype() |
| 594 | |
| 595 | MXFP8_BLOCK_SIZE = 32 |
| 596 | |
| 597 | shape = x.shape[:-1] + (x.shape[-1] // MXFP8_BLOCK_SIZE,) |
| 598 | x_s = torch.zeros(shape, device=x.device, dtype=torch.float32) |
| 599 | |
| 600 | return x.to(dtype), x_s.to(torch.float8_e8m0fnu) |
| 601 | |
| 602 | |
| 603 | def _xpu_mxfp4_quantize_impl( |
| 604 | x: torch.Tensor, |
| 605 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 606 | MXFP4_BLOCK_SIZE = 32 |
| 607 | eps = 1e-10 |
| 608 | assert x.ndim == 2, "input must be 2-D" |
| 609 | assert x.shape[-1] % MXFP4_BLOCK_SIZE == 0, ( |
| 610 | f"last dimension {x.shape[-1]} must be divisible by group_size " |
| 611 | f"{MXFP4_BLOCK_SIZE}" |
| 612 | ) |
| 613 | assert x.is_contiguous(), "input groups must be contiguous" |
| 614 | |
| 615 | M, N = x.shape |
| 616 | |
| 617 | # Packed FP4 output: two nibbles per byte |
| 618 | x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8) |
| 619 | x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32) |
| 620 | |
| 621 | torch.ops._C.per_token_group_quant_mxfp4(x, x_q, x_s, MXFP4_BLOCK_SIZE, eps) |
| 622 | |
| 623 | x_q = x_q.view(torch.float4_e2m1fn_x2) |
| 624 | x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format) |
| 625 | return x_q, x_s |
| 626 | |
| 627 | |
| 628 | def _xpu_mxfp4_quantize_fake( |
| 629 | x: torch.Tensor, |
| 630 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 631 | MXFP4_BLOCK_SIZE = 32 |
| 632 | M, N = x.shape |
| 633 | |
| 634 | # Packed FP4 output: two nibbles per byte |
| 635 | x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8) |
| 636 | x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32) |
| 637 | |
| 638 | x_q = x_q.view(torch.float4_e2m1fn_x2) |
| 639 | x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format) |
| 640 | return x_q, x_s |
| 641 | |
| 642 | |
| 643 | def _xpu_fused_input_norm_impl( |
| 644 | x: torch.Tensor, |
| 645 | weight: torch.Tensor | None, |
| 646 | bias: torch.Tensor | None, |
| 647 | visual_dtype: torch.dtype, |
| 648 | ) -> torch.Tensor: |
| 649 | patches, size = x.shape |
| 650 | out = torch.empty( |
| 651 | (patches, size), |
| 652 | dtype=visual_dtype, |
| 653 | device=x.device, |
| 654 | ) |
| 655 | torch.ops._xpu_C.fused_input_norm(out, x.contiguous(), weight, bias) |
| 656 | return out |
| 657 | |
| 658 | |
| 659 | def _xpu_fused_input_norm_fake( |
| 660 | x: torch.Tensor, |
| 661 | weight: torch.Tensor | None, |
| 662 | bias: torch.Tensor | None, |
| 663 | visual_dtype: torch.dtype, |
| 664 | ) -> torch.Tensor: |
| 665 | patches, size = x.shape |
| 666 | out = torch.empty( |
| 667 | (patches, size), |
| 668 | dtype=visual_dtype, |
| 669 | device=x.device, |
| 670 | ) |
| 671 | return out |
| 672 | |
| 673 | |
| 674 | @triton.jit |
| 675 | def _softplus(x): |
| 676 | return tl.where(x <= 20.0, tl.math.log(tl.math.exp(x) + 1.0), x) |
| 677 | |
| 678 | |
| 679 | @triton.jit |
| 680 | def _selective_scan_fwd_kernel( |
| 681 | # Pointers to input tensors |
| 682 | u_ptr, |
| 683 | delta_ptr, |
| 684 | A_ptr, |
| 685 | B_ptr, |
| 686 | C_ptr, |
| 687 | D_ptr, |
| 688 | z_ptr, |
| 689 | delta_bias_ptr, |
| 690 | # Pointers to output tensors (out aliases delta, out_z aliases z) |
| 691 | out_ptr, |
| 692 | out_z_ptr, |
| 693 | # SSM states |
| 694 | ssm_states_ptr, |
| 695 | # Optional pointers |
| 696 | query_start_loc_ptr, |
| 697 | cache_indices_ptr, |
| 698 | has_initial_state_ptr, |
| 699 | # APC pointers |
| 700 | block_idx_first_ptr, |
| 701 | block_idx_last_ptr, |
| 702 | initial_state_idx_ptr, |
| 703 | cu_chunk_seqlen_ptr, |
| 704 | last_chunk_indices_ptr, |
| 705 | # Dimensions |
| 706 | batch: tl.int32, |
| 707 | dim: tl.int32, |
| 708 | seqlen: tl.int32, |
| 709 | dstate: tl.int32, |
| 710 | n_groups: tl.int32, |
| 711 | dim_ngroups_ratio: tl.int32, |
| 712 | # Strides for u (and out, since out = delta which has same layout) |
| 713 | u_batch_stride: tl.int64, |
| 714 | u_d_stride: tl.int64, |
| 715 | # Strides for delta |
| 716 | delta_batch_stride: tl.int64, |
| 717 | delta_d_stride: tl.int64, |
| 718 | # Strides for A |
| 719 | A_d_stride: tl.int64, |
| 720 | A_dstate_stride: tl.int64, |
| 721 | # Strides for B |
| 722 | B_batch_stride: tl.int64, |
| 723 | B_group_stride: tl.int64, |
| 724 | B_dstate_stride: tl.int64, |
| 725 | # Strides for C |
| 726 | C_batch_stride: tl.int64, |
| 727 | C_group_stride: tl.int64, |
| 728 | C_dstate_stride: tl.int64, |
| 729 | # Strides for z |
| 730 | z_batch_stride: tl.int64, |
| 731 | z_d_stride: tl.int64, |
| 732 | # Strides for out |
| 733 | out_batch_stride: tl.int64, |
| 734 | out_d_stride: tl.int64, |
| 735 | # Strides for out_z |
| 736 | out_z_batch_stride: tl.int64, |
| 737 | out_z_d_stride: tl.int64, |
| 738 | # Strides for ssm_states |
| 739 | ssm_batch_stride: tl.int64, |
| 740 | ssm_dim_stride: tl.int64, |
| 741 | ssm_dstate_stride: tl.int64, |
| 742 | # Cache strides |
| 743 | cache_indices_stride: tl.int64, |
| 744 | # Scalar params |
| 745 | null_block_id: tl.int64, |
| 746 | block_size: tl.int32, |
| 747 | # Compile-time constants |
| 748 | delta_softplus: tl.constexpr, |
| 749 | HAS_D: tl.constexpr, |
| 750 | HAS_Z: tl.constexpr, |
| 751 | HAS_DELTA_BIAS: tl.constexpr, |
| 752 | IS_VARLEN: tl.constexpr, |
| 753 | HAS_CACHE_INDICES: tl.constexpr, |
| 754 | CACHE_ENABLED: tl.constexpr, |
| 755 | BLOCK_DSTATE: tl.constexpr, |
| 756 | ): |
| 757 | batch_idx = tl.program_id(0) |
| 758 | dim_idx = tl.program_id(1) |
| 759 | group_idx = dim_idx // dim_ngroups_ratio |
| 760 | |
| 761 | # Determine sequence boundaries |
| 762 | if IS_VARLEN: |
| 763 | seq_start = tl.load(query_start_loc_ptr + batch_idx).to(tl.int32) |
| 764 | seq_end = tl.load(query_start_loc_ptr + batch_idx + 1).to(tl.int32) |
| 765 | actual_seqlen = seq_end - seq_start |
| 766 | else: |
| 767 | seq_start = 0 |
| 768 | actual_seqlen = seqlen |
| 769 | |
| 770 | # Determine cache index for ssm_states |
| 771 | if CACHE_ENABLED: |
| 772 | init_state_idx = tl.load(initial_state_idx_ptr + batch_idx).to(tl.int32) |
| 773 | load_cache_slot = tl.load( |
| 774 | cache_indices_ptr + batch_idx * cache_indices_stride + init_state_idx |
| 775 | ).to(tl.int64) |
| 776 | if load_cache_slot == null_block_id: |
| 777 | return |
| 778 | elif HAS_CACHE_INDICES: |
| 779 | cache_index = tl.load(cache_indices_ptr + batch_idx).to(tl.int64) |
| 780 | if cache_index == null_block_id: |
| 781 | return |
| 782 | load_cache_slot = cache_index |
| 783 | else: |
| 784 | load_cache_slot = batch_idx.to(tl.int64) |
| 785 | |
| 786 | # Load D value |
| 787 | D_val = 0.0 |
| 788 | if HAS_D: |
| 789 | D_val = tl.load(D_ptr + dim_idx).to(tl.float32) |
| 790 | |
| 791 | # Load delta_bias value |
| 792 | delta_bias_val = 0.0 |
| 793 | if HAS_DELTA_BIAS: |
| 794 | delta_bias_val = tl.load(delta_bias_ptr + dim_idx).to(tl.float32) |
| 795 | |
| 796 | # Load A values for this dim - shape (dstate,) |
| 797 | dstate_offs = tl.arange(0, BLOCK_DSTATE) |
| 798 | dstate_mask = dstate_offs < dstate |
| 799 | A_vals = tl.load( |
| 800 | A_ptr + dim_idx * A_d_stride + dstate_offs * A_dstate_stride, |
| 801 | mask=dstate_mask, |
| 802 | other=0.0, |
| 803 | ).to(tl.float32) |
| 804 | |
| 805 | # Initialize state vector |
| 806 | state = tl.zeros((BLOCK_DSTATE,), dtype=tl.float32) |
| 807 | |
| 808 | # Load initial state if available |
| 809 | has_init = False |
| 810 | if has_initial_state_ptr is not None: |
| 811 | has_init = tl.load(has_initial_state_ptr + batch_idx) |
| 812 | if has_init: |
| 813 | state = tl.load( |
| 814 | ssm_states_ptr |
| 815 | + load_cache_slot * ssm_batch_stride |
| 816 | + dim_idx * ssm_dim_stride |
| 817 | + dstate_offs * ssm_dstate_stride, |
| 818 | mask=dstate_mask, |
| 819 | other=0.0, |
| 820 | ).to(tl.float32) |
| 821 | |
| 822 | # Compute base addresses for u and delta |
| 823 | if IS_VARLEN: |
| 824 | u_base = u_ptr + dim_idx * u_d_stride + seq_start * u_batch_stride |
| 825 | delta_base = ( |
| 826 | delta_ptr + dim_idx * delta_d_stride + seq_start * delta_batch_stride |
| 827 | ) |
| 828 | out_base = out_ptr + dim_idx * out_d_stride + seq_start * out_batch_stride |
| 829 | B_base = B_ptr + group_idx * B_group_stride + seq_start * B_batch_stride |
| 830 | C_base = C_ptr + group_idx * C_group_stride + seq_start * C_batch_stride |
| 831 | else: |
| 832 | u_base = u_ptr + batch_idx * u_batch_stride + dim_idx * u_d_stride |
| 833 | delta_base = ( |
| 834 | delta_ptr + batch_idx * delta_batch_stride + dim_idx * delta_d_stride |
| 835 | ) |
| 836 | out_base = out_ptr + batch_idx * out_batch_stride + dim_idx * out_d_stride |
| 837 | B_base = B_ptr + batch_idx * B_batch_stride + group_idx * B_group_stride |
| 838 | C_base = C_ptr + batch_idx * C_batch_stride + group_idx * C_group_stride |
| 839 | |
| 840 | if HAS_Z: |
| 841 | if IS_VARLEN: |
| 842 | z_base = z_ptr + dim_idx * z_d_stride + seq_start * z_batch_stride |
| 843 | out_z_base = ( |
| 844 | out_z_ptr + dim_idx * out_z_d_stride + seq_start * out_z_batch_stride |
| 845 | ) |
| 846 | else: |
| 847 | z_base = z_ptr + batch_idx * z_batch_stride + dim_idx * z_d_stride |
| 848 | out_z_base = ( |
| 849 | out_z_ptr + batch_idx * out_z_batch_stride + dim_idx * out_z_d_stride |
| 850 | ) |
| 851 | |
| 852 | # Determine chunk boundaries for APC mode |
| 853 | if CACHE_ENABLED: |
| 854 | last_chunk_idx = tl.load(last_chunk_indices_ptr + batch_idx).to(tl.int32) |
| 855 | if batch_idx == 0: |
| 856 | first_chunk_idx = 0 |
| 857 | else: |
| 858 | first_chunk_idx = ( |
| 859 | tl.load(last_chunk_indices_ptr + batch_idx - 1).to(tl.int32) + 1 |
| 860 | ) |
| 861 | n_chunks = last_chunk_idx - first_chunk_idx + 1 |
| 862 | first_chunk_tokens = tl.load(cu_chunk_seqlen_ptr + first_chunk_idx + 1).to( |
| 863 | tl.int32 |
| 864 | ) - tl.load(cu_chunk_seqlen_ptr + first_chunk_idx).to(tl.int32) |
| 865 | block_idx_first = tl.load(block_idx_first_ptr + batch_idx).to(tl.int32) |
| 866 | chunk_start_offset = 0 |
| 867 | if n_chunks > 1 and first_chunk_tokens < block_size: |
| 868 | chunk_start_offset = block_size - first_chunk_tokens |
| 869 | current_position = block_idx_first * block_size + chunk_start_offset |
| 870 | else: |
| 871 | n_chunks = 1 |
| 872 | first_chunk_idx = 0 |
| 873 | |
| 874 | # Sequential scan over the sequence |
| 875 | tokens_processed = 0 |
| 876 | for chunk in range(0, n_chunks if CACHE_ENABLED else 1): |
| 877 | if CACHE_ENABLED: |
| 878 | chunk_tokens = tl.load( |
| 879 | cu_chunk_seqlen_ptr + first_chunk_idx + chunk + 1 |
| 880 | ).to(tl.int32) - tl.load(cu_chunk_seqlen_ptr + first_chunk_idx + chunk).to( |
| 881 | tl.int32 |
| 882 | ) |
| 883 | else: |
| 884 | chunk_tokens = actual_seqlen |
| 885 | |
| 886 | for local_pos in range(chunk_tokens): |
| 887 | pos = tokens_processed + local_pos |
| 888 | # Load u value |
| 889 | u_val = tl.load(u_base + pos).to(tl.float32) |
| 890 | |
| 891 | # Load delta value |
| 892 | delta_val = tl.load(delta_base + pos).to(tl.float32) |
| 893 | |
| 894 | # Apply delta bias |
| 895 | if HAS_DELTA_BIAS: |
| 896 | delta_val = delta_val + delta_bias_val |
| 897 | |
| 898 | # Apply softplus |
| 899 | if delta_softplus: |
| 900 | delta_val = _softplus(delta_val) |
| 901 | |
| 902 | delta_u = delta_val * u_val |
| 903 | |
| 904 | # Compute dA = exp(delta * A) for all dstate elements |
| 905 | dA = tl.exp(delta_val * A_vals) |
| 906 | |
| 907 | # Load B values for this position |
| 908 | B_vals = tl.load( |
| 909 | B_base + dstate_offs * B_dstate_stride + pos, |
| 910 | mask=dstate_mask, |
| 911 | other=0.0, |
| 912 | ).to(tl.float32) |
| 913 | |
| 914 | # Load C values for this position |
| 915 | C_vals = tl.load( |
| 916 | C_base + dstate_offs * C_dstate_stride + pos, |
| 917 | mask=dstate_mask, |
| 918 | other=0.0, |
| 919 | ).to(tl.float32) |
| 920 | |
| 921 | # Update state: state = dA * state + delta * u * B |
| 922 | state = dA * state + delta_u * B_vals |
| 923 | |
| 924 | # Compute output: out = sum(state * C) + D * u |
| 925 | out_val = tl.sum(state * C_vals, axis=0) |
| 926 | if HAS_D: |
| 927 | out_val = out_val + D_val * u_val |
| 928 | |
| 929 | # Store output |
| 930 | tl.store(out_base + pos, out_val.to(out_ptr.dtype.element_ty)) |
| 931 | |
| 932 | if HAS_Z: |
| 933 | z_val = tl.load(z_base + pos).to(tl.float32) |
| 934 | out_z_val = out_val * z_val / (1.0 + tl.exp(-z_val)) |
| 935 | tl.store( |
| 936 | out_z_base + pos, |
| 937 | out_z_val.to(out_z_ptr.dtype.element_ty), |
| 938 | ) |
| 939 | |
| 940 | tokens_processed += chunk_tokens |
| 941 | |
| 942 | # Store intermediate state for APC mode |
| 943 | if CACHE_ENABLED: |
| 944 | if chunk == n_chunks - 1: |
| 945 | store_slot = tl.load( |
| 946 | cache_indices_ptr |
| 947 | + batch_idx * cache_indices_stride |
| 948 | + tl.load(block_idx_last_ptr + batch_idx).to(tl.int32) |
| 949 | ).to(tl.int64) |
| 950 | else: |
| 951 | block_idx_done = (current_position + chunk_tokens - 1) // block_size |
| 952 | store_slot = tl.load( |
| 953 | cache_indices_ptr |
| 954 | + batch_idx * cache_indices_stride |
| 955 | + block_idx_done |
| 956 | ).to(tl.int64) |
| 957 | |
| 958 | tl.store( |
| 959 | ssm_states_ptr |
| 960 | + store_slot * ssm_batch_stride |
| 961 | + dim_idx * ssm_dim_stride |
| 962 | + dstate_offs * ssm_dstate_stride, |
| 963 | state.to(ssm_states_ptr.dtype.element_ty), |
| 964 | mask=dstate_mask, |
| 965 | ) |
| 966 | current_position += chunk_tokens |
| 967 | |
| 968 | # Store final state for non-APC mode |
| 969 | if not CACHE_ENABLED: |
| 970 | tl.store( |
| 971 | ssm_states_ptr |
| 972 | + load_cache_slot * ssm_batch_stride |
| 973 | + dim_idx * ssm_dim_stride |
| 974 | + dstate_offs * ssm_dstate_stride, |
| 975 | state.to(ssm_states_ptr.dtype.element_ty), |
| 976 | mask=dstate_mask, |
| 977 | ) |
| 978 | |
| 979 | |
| 980 | # Global flag to ensure ops are registered only once |
| 981 | _OPS_REGISTERED = False |
| 982 | |
| 983 | |
| 984 | class xpu_ops: |
| 985 | @staticmethod |
| 986 | @torch.compile |
| 987 | def dynamic_per_token_int8_quant_ref( |
| 988 | input: torch.Tensor, use_sym_quant: bool, bits: int |
| 989 | ): |
| 990 | original_sizes = input.size() |
| 991 | # view is not safe in torch.compile if input is not contiguous |
| 992 | input = input.reshape( |
| 993 | -1, original_sizes[-1] |
| 994 | ) # Flatten except for the last dimension |
| 995 | qmin = -(2 ** (bits - 1)) if use_sym_quant else 0 |
| 996 | qmax = 2 ** (bits - 1) - 1 if use_sym_quant else 2**bits - 1 |
| 997 | min_val = torch.min(input, dim=-1)[0].to(dtype=torch.float32).unsqueeze(-1) |
| 998 | max_val = torch.max(input, dim=-1)[0].to(dtype=torch.float32).unsqueeze(-1) |
| 999 | if use_sym_quant: |
| 1000 | scale = ( |
| 1001 | torch.maximum(torch.abs(min_val), torch.abs(max_val)) / qmax |
| 1002 | ).clamp(min=1e-5) |
| 1003 | zero_point = torch.zeros_like(scale).to(dtype=torch.int32) |
| 1004 | else: |
| 1005 | scale = ((max_val - min_val) / qmax).clamp(min=1e-5) |
| 1006 | zero_point = -1 * torch.round(min_val / scale).to(dtype=torch.int32) |
| 1007 | scale = scale.to(dtype=input.dtype) |
| 1008 | quantized = torch.clamp( |
| 1009 | torch.round(input / scale.to(dtype=torch.float32) + zero_point), |
| 1010 | qmin, |
| 1011 | qmax, |
| 1012 | ).to(dtype=torch.int8 if use_sym_quant else torch.uint8) |
| 1013 | return ( |
| 1014 | quantized.view(original_sizes), |
| 1015 | scale.view(original_sizes[:-1] + (1,)), |
| 1016 | zero_point.view(original_sizes[:-1] + (1,)), |
| 1017 | ) |
| 1018 | |
| 1019 | @staticmethod |
| 1020 | def flash_attn_varlen_func( |
| 1021 | q: torch.Tensor, |
| 1022 | k: torch.Tensor, |
| 1023 | v: torch.Tensor, |
| 1024 | cu_seqlens_q: torch.Tensor, |
| 1025 | max_seqlen_q: int, |
| 1026 | max_seqlen_k: int, |
| 1027 | softmax_scale: float | None = None, |
| 1028 | causal: bool = False, |
| 1029 | out: torch.Tensor | None = None, |
| 1030 | block_table: torch.Tensor | None = None, |
| 1031 | alibi_slopes: torch.Tensor | None = None, |
| 1032 | window_size: list[int] | None = None, |
| 1033 | softcap: float | None = 0.0, |
| 1034 | seqused_k: torch.Tensor | None = None, |
| 1035 | cu_seqlens_k: torch.Tensor | None = None, |
| 1036 | # passed in qwen vl |
| 1037 | dropout_p: float = 0.0, |
| 1038 | # The following parameters are not used in xpu kernel currently, |
| 1039 | # we keep API compatible to CUDA's. |
| 1040 | scheduler_metadata=None, |
| 1041 | fa_version: int = 2, |
| 1042 | q_descale=None, |
| 1043 | k_descale=None, |
| 1044 | v_descale=None, |
| 1045 | num_splits=0, |
| 1046 | return_softmax_lse: bool | None = False, |
| 1047 | s_aux: torch.Tensor | None = None, |
| 1048 | return_attn_probs: bool | None = False, |
| 1049 | dynamic_causal: torch.Tensor | None = None, |
| 1050 | mask_mod: Callable | None = None, |
| 1051 | aux_tensors: list | None = None, |
| 1052 | **kwargs, |
| 1053 | ): |
| 1054 | assert cu_seqlens_k is not None or seqused_k is not None, ( |
| 1055 | "cu_seqlens_k or seqused_k must be provided" |
| 1056 | ) |
| 1057 | assert cu_seqlens_k is None or seqused_k is None, ( |
| 1058 | "cu_seqlens_k and seqused_k cannot be provided at the same time" |
| 1059 | ) |
| 1060 | assert block_table is None or seqused_k is not None, ( |
| 1061 | "when enable block_table, seqused_k is needed" |
| 1062 | ) |
| 1063 | assert block_table is not None or cu_seqlens_k is not None, ( |
| 1064 | "when block_table is disabled, cu_seqlens_k is needed" |
| 1065 | ) |
| 1066 | if out is None: |
| 1067 | out = torch.empty(q.shape, dtype=q.dtype, device=q.device) |
| 1068 | real_window_size: tuple[int, int] |
| 1069 | if window_size is None: |
| 1070 | real_window_size = (-1, -1) |
| 1071 | else: |
| 1072 | assert len(window_size) == 2 |
| 1073 | real_window_size = (window_size[0], window_size[1]) # noqa: F841 |
| 1074 | |
| 1075 | return flash_attn_varlen_func( |
| 1076 | out=out, |
| 1077 | q=q, |
| 1078 | k=k, |
| 1079 | v=v, |
| 1080 | cu_seqlens_q=cu_seqlens_q, |
| 1081 | cu_seqlens_k=cu_seqlens_k, |
| 1082 | seqused_k=seqused_k, |
| 1083 | max_seqlen_q=max_seqlen_q, |
| 1084 | max_seqlen_k=max_seqlen_k, |
| 1085 | softmax_scale=softmax_scale, |
| 1086 | causal=causal, |
| 1087 | block_table=block_table, |
| 1088 | s_aux=s_aux, |
| 1089 | window_size=real_window_size, |
| 1090 | # alibi_slopes = alibi_slopes, |
| 1091 | # softcap=softcap, |
| 1092 | return_softmax_lse=return_softmax_lse, |
| 1093 | q_descale=q_descale, |
| 1094 | k_descale=k_descale, |
| 1095 | v_descale=v_descale, |
| 1096 | ) |
| 1097 | |
| 1098 | @staticmethod |
| 1099 | def get_scheduler_metadata( |
| 1100 | batch_size, |
| 1101 | max_seqlen_q, |
| 1102 | max_seqlen_k, |
| 1103 | num_heads_q, |
| 1104 | num_heads_kv, |
| 1105 | headdim, |
| 1106 | cache_seqlens: torch.Tensor, |
| 1107 | qkv_dtype=torch.bfloat16, |
| 1108 | headdim_v=None, |
| 1109 | cu_seqlens_q: torch.Tensor | None = None, |
| 1110 | cu_seqlens_k_new: torch.Tensor | None = None, |
| 1111 | cache_leftpad: torch.Tensor | None = None, |
| 1112 | page_size: int | None = None, |
| 1113 | max_seqlen_k_new=0, |
| 1114 | causal=False, |
| 1115 | window_size=(-1, -1), # -1 means infinite context window |
| 1116 | has_softcap=False, |
| 1117 | num_splits=0, # Can be tuned for speed |
| 1118 | pack_gqa=None, # Can be tuned for speed |
| 1119 | sm_margin=0, # Can be tuned if some SMs are used for communication |
| 1120 | ) -> None: |
| 1121 | logger.warning_once( |
| 1122 | "get_scheduler_metadata is not implemented for xpu_ops, returning None." |
| 1123 | ) |
| 1124 | return None |
| 1125 | |
| 1126 | @staticmethod |
| 1127 | def selective_scan_fwd( |
| 1128 | u: torch.Tensor, |
| 1129 | delta: torch.Tensor, |
| 1130 | A: torch.Tensor, |
| 1131 | B: torch.Tensor, |
| 1132 | C: torch.Tensor, |
| 1133 | D_: torch.Tensor | None, |
| 1134 | z_: torch.Tensor | None, |
| 1135 | delta_bias_: torch.Tensor | None, |
| 1136 | delta_softplus: bool, |
| 1137 | query_start_loc: torch.Tensor | None, |
| 1138 | cache_indices: torch.Tensor | None, |
| 1139 | has_initial_state: torch.Tensor | None, |
| 1140 | ssm_states: torch.Tensor, |
| 1141 | null_block_id: int, |
| 1142 | block_size: int = 1024, |
| 1143 | block_idx_first_scheduled_token: torch.Tensor | None = None, |
| 1144 | block_idx_last_scheduled_token: torch.Tensor | None = None, |
| 1145 | initial_state_idx: torch.Tensor | None = None, |
| 1146 | cu_chunk_seqlen: torch.Tensor | None = None, |
| 1147 | last_chunk_indices: torch.Tensor | None = None, |
| 1148 | ) -> None: |
| 1149 | varlen = query_start_loc is not None |
| 1150 | batch_size = ( |
| 1151 | (query_start_loc.shape[0] - 1) |
| 1152 | if query_start_loc is not None |
| 1153 | else u.shape[0] |
| 1154 | ) |
| 1155 | dim = u.shape[0] if varlen else u.shape[1] |
| 1156 | total_seqlen = u.shape[1] if varlen else u.shape[2] |
| 1157 | dstate = A.size(1) |
| 1158 | n_groups = B.size(0) if varlen else B.size(1) |
| 1159 | dim_ngroups_ratio = dim // n_groups |
| 1160 | |
| 1161 | has_z = z_ is not None |
| 1162 | has_D = D_ is not None |
| 1163 | has_delta_bias = delta_bias_ is not None |
| 1164 | has_cache_indices = cache_indices is not None |
| 1165 | cache_enabled = block_idx_first_scheduled_token is not None |
| 1166 | |
| 1167 | # out and out_z alias delta and z respectively |
| 1168 | out = delta |
| 1169 | out_z = z_ if z_ is not None else delta # won't be used if not has_z |
| 1170 | |
| 1171 | BLOCK_DSTATE = triton.next_power_of_2(dstate) |
| 1172 | |
| 1173 | # Compute strides |
| 1174 | if varlen: |
| 1175 | u_batch_stride = u.stride(1) |
| 1176 | u_d_stride = u.stride(0) |
| 1177 | delta_batch_stride = delta.stride(1) |
| 1178 | delta_d_stride = delta.stride(0) |
| 1179 | B_batch_stride = B.stride(2) |
| 1180 | B_group_stride = B.stride(0) |
| 1181 | B_dstate_stride = B.stride(1) |
| 1182 | C_batch_stride = C.stride(2) |
| 1183 | C_group_stride = C.stride(0) |
| 1184 | C_dstate_stride = C.stride(1) |
| 1185 | out_batch_stride = out.stride(1) |
| 1186 | out_d_stride = out.stride(0) |
| 1187 | if z_ is not None: |
| 1188 | z_batch_stride = z_.stride(1) |
| 1189 | z_d_stride = z_.stride(0) |
| 1190 | out_z_batch_stride = out_z.stride(1) |
| 1191 | out_z_d_stride = out_z.stride(0) |
| 1192 | else: |
| 1193 | z_batch_stride = 0 |
| 1194 | z_d_stride = 0 |
| 1195 | out_z_batch_stride = 0 |
| 1196 | out_z_d_stride = 0 |
| 1197 | else: |
| 1198 | u_batch_stride = u.stride(0) |
| 1199 | u_d_stride = u.stride(1) |
| 1200 | delta_batch_stride = delta.stride(0) |
| 1201 | delta_d_stride = delta.stride(1) |
| 1202 | B_batch_stride = B.stride(0) |
| 1203 | B_group_stride = B.stride(1) |
| 1204 | B_dstate_stride = B.stride(2) |
| 1205 | C_batch_stride = C.stride(0) |
| 1206 | C_group_stride = C.stride(1) |
| 1207 | C_dstate_stride = C.stride(2) |
| 1208 | out_batch_stride = out.stride(0) |
| 1209 | out_d_stride = out.stride(1) |
| 1210 | if z_ is not None: |
| 1211 | z_batch_stride = z_.stride(0) |
| 1212 | z_d_stride = z_.stride(1) |
| 1213 | out_z_batch_stride = out_z.stride(0) |
| 1214 | out_z_d_stride = out_z.stride(1) |
| 1215 | else: |
| 1216 | z_batch_stride = 0 |
| 1217 | z_d_stride = 0 |
| 1218 | out_z_batch_stride = 0 |
| 1219 | out_z_d_stride = 0 |
| 1220 | |
| 1221 | ssm_batch_stride = ssm_states.stride(0) |
| 1222 | ssm_dim_stride = ssm_states.stride(1) |
| 1223 | ssm_dstate_stride = ssm_states.stride(2) |
| 1224 | cache_indices_stride = ( |
| 1225 | cache_indices.stride(0) if cache_indices is not None else 0 |
| 1226 | ) |
| 1227 | |
| 1228 | grid = (batch_size, dim) |
| 1229 | _selective_scan_fwd_kernel[grid]( |
| 1230 | u, |
| 1231 | delta, |
| 1232 | A, |
| 1233 | B, |
| 1234 | C, |
| 1235 | D_ if has_D else u, # dummy, won't be dereferenced |
| 1236 | z_ if has_z else u, # dummy |
| 1237 | delta_bias_ if has_delta_bias else u, # dummy |
| 1238 | out, |
| 1239 | out_z, |
| 1240 | ssm_states, |
| 1241 | query_start_loc if varlen else u, # dummy |
| 1242 | cache_indices if has_cache_indices else u, # dummy |
| 1243 | has_initial_state, |
| 1244 | # APC pointers |
| 1245 | block_idx_first_scheduled_token if cache_enabled else u, |
| 1246 | block_idx_last_scheduled_token if cache_enabled else u, |
| 1247 | initial_state_idx if cache_enabled else u, |
| 1248 | cu_chunk_seqlen if cache_enabled else u, |
| 1249 | last_chunk_indices if cache_enabled else u, |
| 1250 | # Dimensions |
| 1251 | batch_size, |
| 1252 | dim, |
| 1253 | total_seqlen, |
| 1254 | dstate, |
| 1255 | n_groups, |
| 1256 | dim_ngroups_ratio, |
| 1257 | # Strides |
| 1258 | u_batch_stride, |
| 1259 | u_d_stride, |
| 1260 | delta_batch_stride, |
| 1261 | delta_d_stride, |
| 1262 | A.stride(0), |
| 1263 | A.stride(1), |
| 1264 | B_batch_stride, |
| 1265 | B_group_stride, |
| 1266 | B_dstate_stride, |
| 1267 | C_batch_stride, |
| 1268 | C_group_stride, |
| 1269 | C_dstate_stride, |
| 1270 | z_batch_stride, |
| 1271 | z_d_stride, |
| 1272 | out_batch_stride, |
| 1273 | out_d_stride, |
| 1274 | out_z_batch_stride, |
| 1275 | out_z_d_stride, |
| 1276 | ssm_batch_stride, |
| 1277 | ssm_dim_stride, |
| 1278 | ssm_dstate_stride, |
| 1279 | cache_indices_stride, |
| 1280 | null_block_id, |
| 1281 | block_size, |
| 1282 | # Compile-time constants |
| 1283 | delta_softplus=delta_softplus, |
| 1284 | HAS_D=has_D, |
| 1285 | HAS_Z=has_z, |
| 1286 | HAS_DELTA_BIAS=has_delta_bias, |
| 1287 | IS_VARLEN=varlen, |
| 1288 | HAS_CACHE_INDICES=has_cache_indices, |
| 1289 | CACHE_ENABLED=cache_enabled, |
| 1290 | BLOCK_DSTATE=BLOCK_DSTATE, |
| 1291 | ) |
| 1292 | |
| 1293 | @staticmethod |
| 1294 | def register_ops_once() -> None: |
| 1295 | global _OPS_REGISTERED |
| 1296 | if not _OPS_REGISTERED: |
| 1297 | # register all the custom ops here |
| 1298 | if hasattr(torch.ops._C, "gemma_rms_norm"): |
| 1299 | direct_register_custom_op( |
| 1300 | op_name="xpu_gemma_rms_norm", |
| 1301 | op_func=_gemma_rms_norm_impl, |
| 1302 | mutates_args=["out"], |
| 1303 | ) |
| 1304 | |
| 1305 | direct_register_custom_op( |
| 1306 | op_name="xpu_fused_add_gemma_rms_norm", |
| 1307 | op_func=_fused_add_gemma_rms_norm_impl, |
| 1308 | mutates_args=["input", "residual"], |
| 1309 | ) |
| 1310 | |
| 1311 | if hasattr(torch.ops._C, "layer_norm"): |
| 1312 | direct_register_custom_op( |
| 1313 | op_name="xpu_layer_norm", |
| 1314 | op_func=_layer_norm_impl, |
| 1315 | mutates_args=["out"], |
| 1316 | ) |
| 1317 | |
| 1318 | direct_register_custom_op( |
| 1319 | op_name="xpu_fused_add_layer_norm", |
| 1320 | op_func=_fused_add_layer_norm_impl, |
| 1321 | mutates_args=["input", "residual"], |
| 1322 | ) |
| 1323 | |
| 1324 | if hasattr(torch.ops._C, "nemotron_layer_norm"): |
| 1325 | direct_register_custom_op( |
| 1326 | op_name="xpu_nemotron_layer_norm", |
| 1327 | op_func=_nemotron_layer_norm_impl, |
| 1328 | mutates_args=["out"], |
| 1329 | ) |
| 1330 | |
| 1331 | direct_register_custom_op( |
| 1332 | op_name="xpu_fused_add_nemotron_layer_norm", |
| 1333 | op_func=_fused_add_nemotron_layer_norm_impl, |
| 1334 | mutates_args=["input", "residual"], |
| 1335 | ) |
| 1336 | |
| 1337 | direct_register_custom_op( |
| 1338 | op_name="xpu_ops_deepseek_scaling_rope", |
| 1339 | op_func=_xpu_ops_deepseek_scaling_rope_impl, |
| 1340 | mutates_args=[], |
| 1341 | fake_impl=_xpu_ops_deepseek_scaling_rope_fake, |
| 1342 | dispatch_key=current_platform.dispatch_key, |
| 1343 | ) |
| 1344 | |
| 1345 | direct_register_custom_op( |
| 1346 | op_name="xpu_apply_rotary_emb", |
| 1347 | op_func=_xpu_apply_rotary_emb_impl, |
| 1348 | fake_impl=_xpu_apply_rotary_emb_fake, |
| 1349 | ) |
| 1350 | |
| 1351 | direct_register_custom_op( |
| 1352 | op_name="xpu_mxfp8_quantize", |
| 1353 | op_func=_xpu_mxfp8_quantize_impl, |
| 1354 | fake_impl=_xpu_mxfp8_quantize_fake, |
| 1355 | ) |
| 1356 | |
| 1357 | direct_register_custom_op( |
| 1358 | op_name="xpu_mxfp4_quantize", |
| 1359 | op_func=_xpu_mxfp4_quantize_impl, |
| 1360 | fake_impl=_xpu_mxfp4_quantize_fake, |
| 1361 | ) |
| 1362 | |
| 1363 | direct_register_custom_op( |
| 1364 | op_name="xpu_fp8_bmm", |
| 1365 | op_func=_xpu_fp8_bmm_impl, |
| 1366 | fake_impl=_xpu_fp8_bmm_fake, |
| 1367 | ) |
| 1368 | |
| 1369 | direct_register_custom_op( |
| 1370 | op_name="xpu_fp8_mqa_logits", |
| 1371 | op_func=_xpu_fp8_mqa_logits_impl, |
| 1372 | fake_impl=_xpu_fp8_mqa_logits_fake, |
| 1373 | ) |
| 1374 | |
| 1375 | direct_register_custom_op( |
| 1376 | op_name="xpu_fp8_paged_mqa_logits", |
| 1377 | op_func=_xpu_fp8_paged_mqa_logits_impl, |
| 1378 | fake_impl=_xpu_fp8_paged_mqa_logits_fake, |
| 1379 | ) |
| 1380 | |
| 1381 | direct_register_custom_op( |
| 1382 | op_name="gdn_attention_core_xpu", |
| 1383 | op_func=eager_break_during_capture(_gdn_attention_core_xpu_impl), |
| 1384 | mutates_args=["core_attn_out", "z"], |
| 1385 | ) |
| 1386 | |
| 1387 | direct_register_custom_op( |
| 1388 | op_name="xpu_topk_topp_sampler", |
| 1389 | op_func=_topk_topp_sample_impl, |
| 1390 | fake_impl=_topk_topp_sample_fake, |
| 1391 | ) |
| 1392 | |
| 1393 | direct_register_custom_op( |
| 1394 | op_name="xpu_deepseek_fused_indexer_q_rope_fp8", |
| 1395 | op_func=_xpu_deepseek_fused_indexer_q_rope_fp8_impl, |
| 1396 | mutates_args=["index_q_fp8", "index_weights_out"], |
| 1397 | ) |
| 1398 | |
| 1399 | direct_register_custom_op( |
| 1400 | op_name="xpu_deepseek_fused_indexer_q_rope_mxfp4", |
| 1401 | op_func=_xpu_deepseek_fused_indexer_q_rope_mxfp4_impl, |
| 1402 | mutates_args=[ |
| 1403 | "index_q_packed", |
| 1404 | "index_q_scale", |
| 1405 | "index_weights_out", |
| 1406 | ], |
| 1407 | ) |
| 1408 | |
| 1409 | direct_register_custom_op( |
| 1410 | op_name="xpu_fused_input_norm", |
| 1411 | op_func=_xpu_fused_input_norm_impl, |
| 1412 | fake_impl=_xpu_fused_input_norm_fake, |
| 1413 | ) |
| 1414 | |
| 1415 | _OPS_REGISTERED = True |
| 1416 | |
| 1417 | |
| 1418 | xpu_ops.register_ops_once() |