@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f69ede495b
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 06:00 UTC

[XPU][Mamba] Triton-based selective scan forward op for XPU (#43421)

Signed-off-by: Marceli Fylcek <marceli.fylcek@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

Marceli Fylcek committed Jun 2, 2026 at 13:50 UTC f69ede495b3fe97a4b8f6c74d29627f735d46f33
2 files changed +523 -22
vllm/_xpu_ops.py
+474
@@ -8,6 +8,7 @@ from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func
8
9 from vllm.logger import init_logger
10 from vllm.platforms import current_platform
11 +from vllm.triton_utils import tl, triton
12 from vllm.utils.torch_utils import direct_register_custom_op
13
14 logger = init_logger(__name__)
@@ -407,6 +408,312 @@ def _xpu_mxfp4_quantize_fake(
408 return x_q, x_s
409
410
411 +@triton.jit
412 +def _softplus(x):
413 + return tl.where(x <= 20.0, tl.math.log(tl.math.exp(x) + 1.0), x)
414 +
415 +
416 +@triton.jit
417 +def _selective_scan_fwd_kernel(
418 + # Pointers to input tensors
419 + u_ptr,
420 + delta_ptr,
421 + A_ptr,
422 + B_ptr,
423 + C_ptr,
424 + D_ptr,
425 + z_ptr,
426 + delta_bias_ptr,
427 + # Pointers to output tensors (out aliases delta, out_z aliases z)
428 + out_ptr,
429 + out_z_ptr,
430 + # SSM states
431 + ssm_states_ptr,
432 + # Optional pointers
433 + query_start_loc_ptr,
434 + cache_indices_ptr,
435 + has_initial_state_ptr,
436 + # APC pointers
437 + block_idx_first_ptr,
438 + block_idx_last_ptr,
439 + initial_state_idx_ptr,
440 + cu_chunk_seqlen_ptr,
441 + last_chunk_indices_ptr,
442 + # Dimensions
443 + batch: tl.int32,
444 + dim: tl.int32,
445 + seqlen: tl.int32,
446 + dstate: tl.int32,
447 + n_groups: tl.int32,
448 + dim_ngroups_ratio: tl.int32,
449 + # Strides for u (and out, since out = delta which has same layout)
450 + u_batch_stride: tl.int64,
451 + u_d_stride: tl.int64,
452 + # Strides for delta
453 + delta_batch_stride: tl.int64,
454 + delta_d_stride: tl.int64,
455 + # Strides for A
456 + A_d_stride: tl.int64,
457 + A_dstate_stride: tl.int64,
458 + # Strides for B
459 + B_batch_stride: tl.int64,
460 + B_group_stride: tl.int64,
461 + B_dstate_stride: tl.int64,
462 + # Strides for C
463 + C_batch_stride: tl.int64,
464 + C_group_stride: tl.int64,
465 + C_dstate_stride: tl.int64,
466 + # Strides for z
467 + z_batch_stride: tl.int64,
468 + z_d_stride: tl.int64,
469 + # Strides for out
470 + out_batch_stride: tl.int64,
471 + out_d_stride: tl.int64,
472 + # Strides for out_z
473 + out_z_batch_stride: tl.int64,
474 + out_z_d_stride: tl.int64,
475 + # Strides for ssm_states
476 + ssm_batch_stride: tl.int64,
477 + ssm_dim_stride: tl.int64,
478 + ssm_dstate_stride: tl.int64,
479 + # Cache strides
480 + cache_indices_stride: tl.int64,
481 + # Scalar params
482 + null_block_id: tl.int64,
483 + block_size: tl.int32,
484 + # Compile-time constants
485 + delta_softplus: tl.constexpr,
486 + HAS_D: tl.constexpr,
487 + HAS_Z: tl.constexpr,
488 + HAS_DELTA_BIAS: tl.constexpr,
489 + IS_VARLEN: tl.constexpr,
490 + HAS_CACHE_INDICES: tl.constexpr,
491 + CACHE_ENABLED: tl.constexpr,
492 + BLOCK_DSTATE: tl.constexpr,
493 +):
494 + batch_idx = tl.program_id(0)
495 + dim_idx = tl.program_id(1)
496 + group_idx = dim_idx // dim_ngroups_ratio
497 +
498 + # Determine sequence boundaries
499 + if IS_VARLEN:
500 + seq_start = tl.load(query_start_loc_ptr + batch_idx).to(tl.int32)
501 + seq_end = tl.load(query_start_loc_ptr + batch_idx + 1).to(tl.int32)
502 + actual_seqlen = seq_end - seq_start
503 + else:
504 + seq_start = 0
505 + actual_seqlen = seqlen
506 +
507 + # Determine cache index for ssm_states
508 + if CACHE_ENABLED:
509 + init_state_idx = tl.load(initial_state_idx_ptr + batch_idx).to(tl.int32)
510 + load_cache_slot = tl.load(
511 + cache_indices_ptr + batch_idx * cache_indices_stride + init_state_idx
512 + ).to(tl.int64)
513 + if load_cache_slot == null_block_id:
514 + return
515 + elif HAS_CACHE_INDICES:
516 + cache_index = tl.load(cache_indices_ptr + batch_idx).to(tl.int64)
517 + if cache_index == null_block_id:
518 + return
519 + load_cache_slot = cache_index
520 + else:
521 + load_cache_slot = batch_idx.to(tl.int64)
522 +
523 + # Load D value
524 + D_val = 0.0
525 + if HAS_D:
526 + D_val = tl.load(D_ptr + dim_idx).to(tl.float32)
527 +
528 + # Load delta_bias value
529 + delta_bias_val = 0.0
530 + if HAS_DELTA_BIAS:
531 + delta_bias_val = tl.load(delta_bias_ptr + dim_idx).to(tl.float32)
532 +
533 + # Load A values for this dim - shape (dstate,)
534 + dstate_offs = tl.arange(0, BLOCK_DSTATE)
535 + dstate_mask = dstate_offs < dstate
536 + A_vals = tl.load(
537 + A_ptr + dim_idx * A_d_stride + dstate_offs * A_dstate_stride,
538 + mask=dstate_mask,
539 + other=0.0,
540 + ).to(tl.float32)
541 +
542 + # Initialize state vector
543 + state = tl.zeros((BLOCK_DSTATE,), dtype=tl.float32)
544 +
545 + # Load initial state if available
546 + has_init = False
547 + if has_initial_state_ptr is not None:
548 + has_init = tl.load(has_initial_state_ptr + batch_idx)
549 + if has_init:
550 + state = tl.load(
551 + ssm_states_ptr
552 + + load_cache_slot * ssm_batch_stride
553 + + dim_idx * ssm_dim_stride
554 + + dstate_offs * ssm_dstate_stride,
555 + mask=dstate_mask,
556 + other=0.0,
557 + ).to(tl.float32)
558 +
559 + # Compute base addresses for u and delta
560 + if IS_VARLEN:
561 + u_base = u_ptr + dim_idx * u_d_stride + seq_start * u_batch_stride
562 + delta_base = (
563 + delta_ptr + dim_idx * delta_d_stride + seq_start * delta_batch_stride
564 + )
565 + out_base = out_ptr + dim_idx * out_d_stride + seq_start * out_batch_stride
566 + B_base = B_ptr + group_idx * B_group_stride + seq_start * B_batch_stride
567 + C_base = C_ptr + group_idx * C_group_stride + seq_start * C_batch_stride
568 + else:
569 + u_base = u_ptr + batch_idx * u_batch_stride + dim_idx * u_d_stride
570 + delta_base = (
571 + delta_ptr + batch_idx * delta_batch_stride + dim_idx * delta_d_stride
572 + )
573 + out_base = out_ptr + batch_idx * out_batch_stride + dim_idx * out_d_stride
574 + B_base = B_ptr + batch_idx * B_batch_stride + group_idx * B_group_stride
575 + C_base = C_ptr + batch_idx * C_batch_stride + group_idx * C_group_stride
576 +
577 + if HAS_Z:
578 + if IS_VARLEN:
579 + z_base = z_ptr + dim_idx * z_d_stride + seq_start * z_batch_stride
580 + out_z_base = (
581 + out_z_ptr + dim_idx * out_z_d_stride + seq_start * out_z_batch_stride
582 + )
583 + else:
584 + z_base = z_ptr + batch_idx * z_batch_stride + dim_idx * z_d_stride
585 + out_z_base = (
586 + out_z_ptr + batch_idx * out_z_batch_stride + dim_idx * out_z_d_stride
587 + )
588 +
589 + # Determine chunk boundaries for APC mode
590 + if CACHE_ENABLED:
591 + last_chunk_idx = tl.load(last_chunk_indices_ptr + batch_idx).to(tl.int32)
592 + if batch_idx == 0:
593 + first_chunk_idx = 0
594 + else:
595 + first_chunk_idx = (
596 + tl.load(last_chunk_indices_ptr + batch_idx - 1).to(tl.int32) + 1
597 + )
598 + n_chunks = last_chunk_idx - first_chunk_idx + 1
599 + first_chunk_tokens = tl.load(cu_chunk_seqlen_ptr + first_chunk_idx + 1).to(
600 + tl.int32
601 + ) - tl.load(cu_chunk_seqlen_ptr + first_chunk_idx).to(tl.int32)
602 + block_idx_first = tl.load(block_idx_first_ptr + batch_idx).to(tl.int32)
603 + chunk_start_offset = 0
604 + if n_chunks > 1 and first_chunk_tokens < block_size:
605 + chunk_start_offset = block_size - first_chunk_tokens
606 + current_position = block_idx_first * block_size + chunk_start_offset
607 + else:
608 + n_chunks = 1
609 + first_chunk_idx = 0
610 +
611 + # Sequential scan over the sequence
612 + tokens_processed = 0
613 + for chunk in range(0, n_chunks if CACHE_ENABLED else 1):
614 + if CACHE_ENABLED:
615 + chunk_tokens = tl.load(
616 + cu_chunk_seqlen_ptr + first_chunk_idx + chunk + 1
617 + ).to(tl.int32) - tl.load(cu_chunk_seqlen_ptr + first_chunk_idx + chunk).to(
618 + tl.int32
619 + )
620 + else:
621 + chunk_tokens = actual_seqlen
622 +
623 + for local_pos in range(chunk_tokens):
624 + pos = tokens_processed + local_pos
625 + # Load u value
626 + u_val = tl.load(u_base + pos).to(tl.float32)
627 +
628 + # Load delta value
629 + delta_val = tl.load(delta_base + pos).to(tl.float32)
630 +
631 + # Apply delta bias
632 + if HAS_DELTA_BIAS:
633 + delta_val = delta_val + delta_bias_val
634 +
635 + # Apply softplus
636 + if delta_softplus:
637 + delta_val = _softplus(delta_val)
638 +
639 + delta_u = delta_val * u_val
640 +
641 + # Compute dA = exp(delta * A) for all dstate elements
642 + dA = tl.exp(delta_val * A_vals)
643 +
644 + # Load B values for this position
645 + B_vals = tl.load(
646 + B_base + dstate_offs * B_dstate_stride + pos,
647 + mask=dstate_mask,
648 + other=0.0,
649 + ).to(tl.float32)
650 +
651 + # Load C values for this position
652 + C_vals = tl.load(
653 + C_base + dstate_offs * C_dstate_stride + pos,
654 + mask=dstate_mask,
655 + other=0.0,
656 + ).to(tl.float32)
657 +
658 + # Update state: state = dA * state + delta * u * B
659 + state = dA * state + delta_u * B_vals
660 +
661 + # Compute output: out = sum(state * C) + D * u
662 + out_val = tl.sum(state * C_vals, axis=0)
663 + if HAS_D:
664 + out_val = out_val + D_val * u_val
665 +
666 + # Store output
667 + tl.store(out_base + pos, out_val.to(out_ptr.dtype.element_ty))
668 +
669 + if HAS_Z:
670 + z_val = tl.load(z_base + pos).to(tl.float32)
671 + out_z_val = out_val * z_val / (1.0 + tl.exp(-z_val))
672 + tl.store(
673 + out_z_base + pos,
674 + out_z_val.to(out_z_ptr.dtype.element_ty),
675 + )
676 +
677 + tokens_processed += chunk_tokens
678 +
679 + # Store intermediate state for APC mode
680 + if CACHE_ENABLED:
681 + if chunk == n_chunks - 1:
682 + store_slot = tl.load(
683 + cache_indices_ptr
684 + + batch_idx * cache_indices_stride
685 + + tl.load(block_idx_last_ptr + batch_idx).to(tl.int32)
686 + ).to(tl.int64)
687 + else:
688 + block_idx_done = (current_position + chunk_tokens - 1) // block_size
689 + store_slot = tl.load(
690 + cache_indices_ptr
691 + + batch_idx * cache_indices_stride
692 + + block_idx_done
693 + ).to(tl.int64)
694 +
695 + tl.store(
696 + ssm_states_ptr
697 + + store_slot * ssm_batch_stride
698 + + dim_idx * ssm_dim_stride
699 + + dstate_offs * ssm_dstate_stride,
700 + state.to(ssm_states_ptr.dtype.element_ty),
701 + mask=dstate_mask,
702 + )
703 + current_position += chunk_tokens
704 +
705 + # Store final state for non-APC mode
706 + if not CACHE_ENABLED:
707 + tl.store(
708 + ssm_states_ptr
709 + + load_cache_slot * ssm_batch_stride
710 + + dim_idx * ssm_dim_stride
711 + + dstate_offs * ssm_dstate_stride,
712 + state.to(ssm_states_ptr.dtype.element_ty),
713 + mask=dstate_mask,
714 + )
715 +
716 +
717 # Global flag to ensure ops are registered only once
718 _OPS_REGISTERED = False
719
@@ -549,6 +856,173 @@ class xpu_ops:
856 )
857 return None
858
859 + @staticmethod
860 + def selective_scan_fwd(
861 + u: torch.Tensor,
862 + delta: torch.Tensor,
863 + A: torch.Tensor,
864 + B: torch.Tensor,
865 + C: torch.Tensor,
866 + D_: torch.Tensor | None,
867 + z_: torch.Tensor | None,
868 + delta_bias_: torch.Tensor | None,
869 + delta_softplus: bool,
870 + query_start_loc: torch.Tensor | None,
871 + cache_indices: torch.Tensor | None,
872 + has_initial_state: torch.Tensor | None,
873 + ssm_states: torch.Tensor,
874 + null_block_id: int,
875 + block_size: int = 1024,
876 + block_idx_first_scheduled_token: torch.Tensor | None = None,
877 + block_idx_last_scheduled_token: torch.Tensor | None = None,
878 + initial_state_idx: torch.Tensor | None = None,
879 + cu_chunk_seqlen: torch.Tensor | None = None,
880 + last_chunk_indices: torch.Tensor | None = None,
881 + ) -> None:
882 + varlen = query_start_loc is not None
883 + batch_size = (
884 + (query_start_loc.shape[0] - 1)
885 + if query_start_loc is not None
886 + else u.shape[0]
887 + )
888 + dim = u.shape[0] if varlen else u.shape[1]
889 + total_seqlen = u.shape[1] if varlen else u.shape[2]
890 + dstate = A.size(1)
891 + n_groups = B.size(0) if varlen else B.size(1)
892 + dim_ngroups_ratio = dim // n_groups
893 +
894 + has_z = z_ is not None
895 + has_D = D_ is not None
896 + has_delta_bias = delta_bias_ is not None
897 + has_cache_indices = cache_indices is not None
898 + cache_enabled = block_idx_first_scheduled_token is not None
899 +
900 + # out and out_z alias delta and z respectively
901 + out = delta
902 + out_z = z_ if z_ is not None else delta # won't be used if not has_z
903 +
904 + BLOCK_DSTATE = triton.next_power_of_2(dstate)
905 +
906 + # Compute strides
907 + if varlen:
908 + u_batch_stride = u.stride(1)
909 + u_d_stride = u.stride(0)
910 + delta_batch_stride = delta.stride(1)
911 + delta_d_stride = delta.stride(0)
912 + B_batch_stride = B.stride(2)
913 + B_group_stride = B.stride(0)
914 + B_dstate_stride = B.stride(1)
915 + C_batch_stride = C.stride(2)
916 + C_group_stride = C.stride(0)
917 + C_dstate_stride = C.stride(1)
918 + out_batch_stride = out.stride(1)
919 + out_d_stride = out.stride(0)
920 + if z_ is not None:
921 + z_batch_stride = z_.stride(1)
922 + z_d_stride = z_.stride(0)
923 + out_z_batch_stride = out_z.stride(1)
924 + out_z_d_stride = out_z.stride(0)
925 + else:
926 + z_batch_stride = 0
927 + z_d_stride = 0
928 + out_z_batch_stride = 0
929 + out_z_d_stride = 0
930 + else:
931 + u_batch_stride = u.stride(0)
932 + u_d_stride = u.stride(1)
933 + delta_batch_stride = delta.stride(0)
934 + delta_d_stride = delta.stride(1)
935 + B_batch_stride = B.stride(0)
936 + B_group_stride = B.stride(1)
937 + B_dstate_stride = B.stride(2)
938 + C_batch_stride = C.stride(0)
939 + C_group_stride = C.stride(1)
940 + C_dstate_stride = C.stride(2)
941 + out_batch_stride = out.stride(0)
942 + out_d_stride = out.stride(1)
943 + if z_ is not None:
944 + z_batch_stride = z_.stride(0)
945 + z_d_stride = z_.stride(1)
946 + out_z_batch_stride = out_z.stride(0)
947 + out_z_d_stride = out_z.stride(1)
948 + else:
949 + z_batch_stride = 0
950 + z_d_stride = 0
951 + out_z_batch_stride = 0
952 + out_z_d_stride = 0
953 +
954 + ssm_batch_stride = ssm_states.stride(0)
955 + ssm_dim_stride = ssm_states.stride(1)
956 + ssm_dstate_stride = ssm_states.stride(2)
957 + cache_indices_stride = (
958 + cache_indices.stride(0) if cache_indices is not None else 0
959 + )
960 +
961 + grid = (batch_size, dim)
962 + _selective_scan_fwd_kernel[grid](
963 + u,
964 + delta,
965 + A,
966 + B,
967 + C,
968 + D_ if has_D else u, # dummy, won't be dereferenced
969 + z_ if has_z else u, # dummy
970 + delta_bias_ if has_delta_bias else u, # dummy
971 + out,
972 + out_z,
973 + ssm_states,
974 + query_start_loc if varlen else u, # dummy
975 + cache_indices if has_cache_indices else u, # dummy
976 + has_initial_state,
977 + # APC pointers
978 + block_idx_first_scheduled_token if cache_enabled else u,
979 + block_idx_last_scheduled_token if cache_enabled else u,
980 + initial_state_idx if cache_enabled else u,
981 + cu_chunk_seqlen if cache_enabled else u,
982 + last_chunk_indices if cache_enabled else u,
983 + # Dimensions
984 + batch_size,
985 + dim,
986 + total_seqlen,
987 + dstate,
988 + n_groups,
989 + dim_ngroups_ratio,
990 + # Strides
991 + u_batch_stride,
992 + u_d_stride,
993 + delta_batch_stride,
994 + delta_d_stride,
995 + A.stride(0),
996 + A.stride(1),
997 + B_batch_stride,
998 + B_group_stride,
999 + B_dstate_stride,
1000 + C_batch_stride,
1001 + C_group_stride,
1002 + C_dstate_stride,
1003 + z_batch_stride,
1004 + z_d_stride,
1005 + out_batch_stride,
1006 + out_d_stride,
1007 + out_z_batch_stride,
1008 + out_z_d_stride,
1009 + ssm_batch_stride,
1010 + ssm_dim_stride,
1011 + ssm_dstate_stride,
1012 + cache_indices_stride,
1013 + null_block_id,
1014 + block_size,
1015 + # Compile-time constants
1016 + delta_softplus=delta_softplus,
1017 + HAS_D=has_D,
1018 + HAS_Z=has_z,
1019 + HAS_DELTA_BIAS=has_delta_bias,
1020 + IS_VARLEN=varlen,
1021 + HAS_CACHE_INDICES=has_cache_indices,
1022 + CACHE_ENABLED=cache_enabled,
1023 + BLOCK_DSTATE=BLOCK_DSTATE,
1024 + )
1025 +
1026 @staticmethod
1027 def register_ops_once() -> None:
1028 global _OPS_REGISTERED
vllm/model_executor/layers/mamba/ops/mamba_ssm.py
+49 -22
@@ -21,6 +21,9 @@ from vllm.platforms import current_platform
21 from vllm.triton_utils import HAS_TRITON, tl, triton
22 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
23
24 +if current_platform.is_xpu():
25 + from vllm._xpu_ops import xpu_ops
26 +
27 logger = init_logger(__name__)
28
29 TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0"))
@@ -790,28 +793,52 @@ def selective_scan_fn(
793 if C.dim() == 2 and query_start_loc is not None:
794 C = C.unsqueeze(0)
795
793 - ops.selective_scan_fwd(
794 - u,
795 - delta,
796 - A,
797 - B,
798 - C,
799 - D,
800 - z,
801 - delta_bias,
802 - delta_softplus,
803 - query_start_loc,
804 - cache_indices,
805 - has_initial_state,
806 - ssm_states,
807 - null_block_id,
808 - block_size,
809 - block_idx_first_scheduled_token,
810 - block_idx_last_scheduled_token,
811 - initial_state_idx,
812 - cu_chunk_seqlen,
813 - last_chunk_indices,
814 - )
796 + if current_platform.is_xpu():
797 + xpu_ops.selective_scan_fwd(
798 + u,
799 + delta,
800 + A,
801 + B,
802 + C,
803 + D,
804 + z,
805 + delta_bias,
806 + delta_softplus,
807 + query_start_loc,
808 + cache_indices,
809 + has_initial_state,
810 + ssm_states,
811 + null_block_id,
812 + block_size,
813 + block_idx_first_scheduled_token,
814 + block_idx_last_scheduled_token,
815 + initial_state_idx,
816 + cu_chunk_seqlen,
817 + last_chunk_indices,
818 + )
819 + else:
820 + ops.selective_scan_fwd(
821 + u,
822 + delta,
823 + A,
824 + B,
825 + C,
826 + D,
827 + z,
828 + delta_bias,
829 + delta_softplus,
830 + query_start_loc,
831 + cache_indices,
832 + has_initial_state,
833 + ssm_states,
834 + null_block_id,
835 + block_size,
836 + block_idx_first_scheduled_token,
837 + block_idx_last_scheduled_token,
838 + initial_state_idx,
839 + cu_chunk_seqlen,
840 + last_chunk_indices,
841 + )
842
843 if z is None:
844 return delta # output written inplace to delta