@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ae0f69b16a

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 11, 20:00 UTC

Add SpecDec support to `selective_state_update` (#29488)

Signed-off-by: Roi Koren <roik@nvidia.com>

roikoren755 committed Dec 8, 2025 at 23:45 UTC ae0f69b16aaa54197257267a8262555beb7d5ae3

2 files changed +507 -74

tests/kernels/mamba/test_mamba_ssm.py

+325

@@ -425,6 +425,80 @@ def test_selective_state_update(dim, dstate, has_z, itype):

425 assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)

426

427

428 +@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])

429 +@pytest.mark.parametrize("has_z", [False, True])

430 +@pytest.mark.parametrize("dstate", [16, 64])

431 +@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])

432 +@pytest.mark.parametrize("max_seq_len", [1, 2, 4])

433 +def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len):

434 + device = "cuda"

435 + rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)

436 + if itype == torch.bfloat16:

437 + rtol, atol = 5e-2, 1.5e-1

438 + if torch.version.hip:

439 + atol *= 2

440 + # set seed

441 + current_platform.seed_everything(0)

442 + batch_size = 4

443 + token_counts = torch.randint(1, max_seq_len + 1, (batch_size,), device=device)

444 + total_tokens = int(token_counts.sum().item())

445 + cu_seqlens = torch.tensor(

446 + [0] + torch.cumsum(token_counts, dim=0).tolist(),

447 + dtype=torch.int32,

448 + device=device,

449 + )

450 + state = torch.randn(batch_size, dim, dstate, dtype=itype, device=device)

451 + x = torch.randn(total_tokens, dim, device=device, dtype=itype)

452 + out = torch.empty_like(x)

453 + dt = torch.randn(total_tokens, dim, device=device, dtype=itype)

454 + dt_bias = torch.rand(dim, device=device) - 4.0

455 + A = -torch.rand(dim, dstate, device=device) - 1.0

456 + B = torch.randn(total_tokens, dstate, device=device)

457 + C = torch.randn(total_tokens, dstate, device=device)

458 + D = torch.randn(dim, device=device)

459 + z = torch.randn_like(x) if has_z else None

460 + state_ref = state.detach().clone()

461 + selective_state_update(

462 + state,

463 + x,

464 + dt,

465 + A,

466 + B,

467 + C,

468 + D=D,

469 + z=z,

470 + dt_bias=dt_bias,

471 + dt_softplus=True,

472 + out=out,

473 + cu_seqlens=cu_seqlens,

474 + )

475 +

476 + out_ref_list = []

477 + for seq_idx in range(batch_size):

478 + start_idx = cu_seqlens[seq_idx].item()

479 + end_idx = cu_seqlens[seq_idx + 1].item()

480 + num_tokens = end_idx - start_idx

481 + for token_idx in range(num_tokens):

482 + idx = start_idx + token_idx

483 + out_ref_list.append(

484 + selective_state_update_ref(

485 + state_ref[seq_idx : seq_idx + 1],

486 + x[idx : idx + 1],

487 + dt[idx : idx + 1],

488 + A,

489 + B[idx : idx + 1],

490 + C[idx : idx + 1],

491 + D=D,

492 + z=z[idx : idx + 1] if has_z else None,

493 + dt_bias=dt_bias,

494 + dt_softplus=True,

495 + )

496 + )

497 + out_ref = torch.cat(out_ref_list, dim=0)

498 + assert torch.allclose(state, state_ref, rtol=rtol, atol=atol)

499 + assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)

500 +

501 +

502 @pytest.mark.parametrize("wtype", [torch.float32])

503 @pytest.mark.parametrize("itype", [torch.float32])

504 @pytest.mark.parametrize("seqlen", [1, 256, 1024, 4096])

@@ -766,3 +840,254 @@ def test_selective_state_update_with_heads_with_batch_indices(

840 print(f"Output mean diff: {(out - out_ref).abs().mean().item()}")

841 assert torch.allclose(state[state_indices, :], state_ref, rtol=rtol, atol=atol)

842 assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)

843 +

844 +

845 +@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])

846 +@pytest.mark.parametrize("has_z", [False, True])

847 +@pytest.mark.parametrize("dstate", [16, 64])

848 +@pytest.mark.parametrize("dim", [2048, 4096])

849 +@pytest.mark.parametrize("max_seq_len", [2, 4])

850 +def test_selective_state_update_with_num_accepted_tokens(

851 + dim, dstate, has_z, itype, max_seq_len

852 +):

853 + device = "cuda"

854 + rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)

855 + if itype == torch.bfloat16:

856 + rtol, atol = 5e-2, 1.5e-1

857 + if torch.version.hip:

858 + atol *= 2

859 +

860 + current_platform.seed_everything(0)

861 + batch_size = 4

862 +

863 + tokens_per_seq = torch.randint(1, max_seq_len + 1, (batch_size,), device=device)

864 + total_tokens = int(tokens_per_seq.sum().item())

865 +

866 + num_accepted_tokens = torch.randint(0, max_seq_len, (batch_size,), device=device)

867 + num_accepted_tokens[0] = 0 # Add edge-case of no accepted tokens

868 + num_accepted_tokens[1] = max_seq_len # Add edge-case of all tokens accepted

869 +

870 + cu_seqlens = torch.tensor(

871 + [0] + torch.cumsum(tokens_per_seq, dim=0).tolist(),

872 + dtype=torch.int32,

873 + device=device,

874 + )

875 +

876 + total_state_slots = 50

877 + state = torch.randn(total_state_slots, dim, dstate, dtype=itype, device=device)

878 +

879 + state_batch_indices = torch.full(

880 + (batch_size, max_seq_len), PAD_SLOT_ID, dtype=torch.int32, device=device

881 + )

882 + initial_state_slots = torch.randint(

883 + 0, 15, (batch_size,), device=device, dtype=torch.int32

884 + )

885 + for seq_idx in range(batch_size):

886 + token_pos = max(num_accepted_tokens[seq_idx].item() - 1, 0)

887 + state_batch_indices[seq_idx, token_pos] = initial_state_slots[seq_idx]

888 +

889 + dst_state_batch_indices = torch.full(

890 + (batch_size, max_seq_len), PAD_SLOT_ID, dtype=torch.int32, device=device

891 + )

892 + slot_offset = 15

893 + dst_slots_map = {}

894 + for seq_idx in range(batch_size):

895 + for token_idx in range(tokens_per_seq[seq_idx].item()):

896 + dst_state_batch_indices[seq_idx, token_idx] = slot_offset

897 + dst_slots_map[(seq_idx, token_idx)] = slot_offset

898 + slot_offset += 1

899 +

900 + x = torch.randn(total_tokens, dim, device=device, dtype=itype)

901 + out = torch.empty_like(x)

902 + dt = torch.randn(total_tokens, dim, device=device, dtype=itype)

903 + dt_bias = torch.rand(dim, device=device) - 4.0

904 + A = -torch.rand(dim, dstate, device=device) - 1.0

905 + B = torch.randn(total_tokens, dstate, device=device)

906 + C = torch.randn(total_tokens, dstate, device=device)

907 + D = torch.randn(dim, device=device)

908 + z = torch.randn_like(x) if has_z else None

909 +

910 + state_ref_intermediate = {}

911 + out_ref_list = []

912 +

913 + for seq_idx in range(batch_size):

914 + seq_start = cu_seqlens[seq_idx].item()

915 + seq_end = cu_seqlens[seq_idx + 1].item()

916 + num_tokens = seq_end - seq_start

917 +

918 + token_pos = max(num_accepted_tokens[seq_idx].item() - 1, 0)

919 + initial_slot = state_batch_indices[seq_idx, token_pos].item()

920 + state_seq = state[initial_slot : initial_slot + 1].clone()

921 +

922 + for token_idx in range(num_tokens):

923 + global_idx = seq_start + token_idx

924 +

925 + out_token = selective_state_update_ref(

926 + state_seq,

927 + x[global_idx : global_idx + 1],

928 + dt[global_idx : global_idx + 1],

929 + A,

930 + B[global_idx : global_idx + 1],

931 + C[global_idx : global_idx + 1],

932 + D=D,

933 + z=z[global_idx : global_idx + 1] if has_z else None,

934 + dt_bias=dt_bias,

935 + dt_softplus=True,

936 + )

937 + out_ref_list.append(out_token)

938 + state_ref_intermediate[(seq_idx, token_idx)] = state_seq.clone()

939 +

940 + out_ref = torch.cat(out_ref_list, dim=0)

941 +

942 + selective_state_update(

943 + state,

944 + x,

945 + dt,

946 + A,

947 + B,

948 + C,

949 + D=D,

950 + z=z,

951 + dt_bias=dt_bias,

952 + dt_softplus=True,

953 + out=out,

954 + cu_seqlens=cu_seqlens,

955 + state_batch_indices=state_batch_indices,

956 + dst_state_batch_indices=dst_state_batch_indices,

957 + num_accepted_tokens=num_accepted_tokens,

958 + pad_slot_id=PAD_SLOT_ID,

959 + )

960 +

961 + assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)

962 +

963 + for seq_idx in range(batch_size):

964 + num_tokens = tokens_per_seq[seq_idx].item()

965 + for token_idx in range(num_tokens):

966 + dst_slot = dst_slots_map[(seq_idx, token_idx)]

967 + state_ref = state_ref_intermediate[(seq_idx, token_idx)].squeeze(0)

968 + assert torch.allclose(state[dst_slot], state_ref, rtol=rtol, atol=atol)

969 +

970 +

971 +@pytest.mark.parametrize("itype", [torch.float32, torch.bfloat16])

972 +@pytest.mark.parametrize("has_z", [False, True])

973 +@pytest.mark.parametrize("dstate", [16, 64])

974 +@pytest.mark.parametrize("dim", [2048, 4096])

975 +@pytest.mark.parametrize("max_seq_len", [2, 4])

976 +def test_selective_state_update_varlen_with_num_accepted(

977 + dim, dstate, has_z, itype, max_seq_len

978 +):

979 + device = "cuda"

980 + rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2)

981 + if itype == torch.bfloat16:

982 + rtol, atol = 5e-2, 1.5e-1

983 + if torch.version.hip:

984 + atol *= 2

985 +

986 + current_platform.seed_everything(0)

987 + batch_size = 4

988 +

989 + tokens_per_seq = torch.randint(1, max_seq_len + 1, (batch_size,), device=device)

990 + total_tokens = int(tokens_per_seq.sum().item())

991 +

992 + num_accepted_tokens = torch.randint(0, max_seq_len, (batch_size,), device=device)

993 + num_accepted_tokens[0] = 0 # Add edge-case of no accepted tokens

994 + num_accepted_tokens[1] = max_seq_len # Add edge-case of all tokens accepted

995 +

996 + cu_seqlens = torch.tensor(

997 + [0] + torch.cumsum(tokens_per_seq, dim=0).tolist(),

998 + dtype=torch.int32,

999 + device=device,

1000 + )

1001 +

1002 + total_state_slots = 50

1003 + state = torch.randn(total_state_slots, dim, dstate, dtype=itype, device=device)

1004 +

1005 + state_batch_indices = torch.full(

1006 + (batch_size, max_seq_len), PAD_SLOT_ID, dtype=torch.int32, device=device

1007 + )

1008 +

1009 + initial_state_slots = torch.randint(

1010 + 0, 15, (batch_size,), device=device, dtype=torch.int32

1011 + )

1012 + for seq_idx in range(batch_size):

1013 + token_pos = max(num_accepted_tokens[seq_idx].item() - 1, 0)

1014 + state_batch_indices[seq_idx, token_pos] = initial_state_slots[seq_idx]

1015 +

1016 + dst_state_batch_indices = torch.full(

1017 + (batch_size, max_seq_len), PAD_SLOT_ID, dtype=torch.int32, device=device

1018 + )

1019 +

1020 + slot_offset = 15

1021 + dst_slots_map = {}

1022 + for seq_idx in range(batch_size):

1023 + for token_idx in range(tokens_per_seq[seq_idx].item()):

1024 + dst_state_batch_indices[seq_idx, token_idx] = slot_offset

1025 + dst_slots_map[(seq_idx, token_idx)] = slot_offset

1026 + slot_offset += 1

1027 +

1028 + x = torch.randn(total_tokens, dim, device=device, dtype=itype)

1029 + out = torch.empty_like(x)

1030 + dt = torch.randn(total_tokens, dim, device=device, dtype=itype)

1031 + dt_bias = torch.rand(dim, device=device) - 4.0

1032 + A = -torch.rand(dim, dstate, device=device) - 1.0

1033 + B = torch.randn(total_tokens, dstate, device=device)

1034 + C = torch.randn(total_tokens, dstate, device=device)

1035 + D = torch.randn(dim, device=device)

1036 + z = torch.randn_like(x) if has_z else None

1037 +

1038 + state_ref_intermediate = {}

1039 +

1040 + for seq_idx in range(batch_size):

1041 + seq_start = cu_seqlens[seq_idx].item()

1042 + seq_end = cu_seqlens[seq_idx + 1].item()

1043 + num_tokens = seq_end - seq_start

1044 +

1045 + token_pos = max(num_accepted_tokens[seq_idx].item() - 1, 0)

1046 + initial_slot = state_batch_indices[seq_idx, token_pos].item()

1047 + state_seq = state[initial_slot : initial_slot + 1].clone()

1048 +

1049 + for token_idx in range(num_tokens):

1050 + global_idx = seq_start + token_idx

1051 +

1052 + selective_state_update_ref(

1053 + state_seq,

1054 + x[global_idx : global_idx + 1],

1055 + dt[global_idx : global_idx + 1],

1056 + A,

1057 + B[global_idx : global_idx + 1],

1058 + C[global_idx : global_idx + 1],

1059 + D=D,

1060 + z=z[global_idx : global_idx + 1] if has_z else None,

1061 + dt_bias=dt_bias,

1062 + dt_softplus=True,

1063 + )

1064 +

1065 + state_ref_intermediate[(seq_idx, token_idx)] = state_seq.clone()

1066 +

1067 + selective_state_update(

1068 + state,

1069 + x,

1070 + dt,

1071 + A,

1072 + B,

1073 + C,

1074 + D=D,

1075 + z=z,

1076 + dt_bias=dt_bias,

1077 + dt_softplus=True,

1078 + out=out,

1079 + cu_seqlens=cu_seqlens,

1080 + state_batch_indices=state_batch_indices,

1081 + dst_state_batch_indices=dst_state_batch_indices,

1082 + num_accepted_tokens=num_accepted_tokens,

1083 + pad_slot_id=PAD_SLOT_ID,

1084 + )

1085 +

1086 + for seq_idx in range(batch_size):

1087 + num_tokens = tokens_per_seq[seq_idx].item()

1088 +

1089 + for token_idx in range(num_tokens):

1090 + dst_slot = dst_slots_map[(seq_idx, token_idx)]

1091 + state_ref = state_ref_intermediate[(seq_idx, token_idx)].squeeze(0)

1092 +

1093 + assert torch.allclose(state[dst_slot], state_ref, rtol=rtol, atol=atol)

vllm/model_executor/layers/mamba/ops/mamba_ssm.py

+182 -74

@@ -36,10 +36,14 @@ else:

36 is not None

37 }

38 )

39 +@triton.heuristics(

40 + {"IS_SPEC_DECODING": lambda args: args["num_accepted_tokens_ptr"] is not None}

41 +)

42 +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens_ptr"] is not None})

43 @triton.heuristics(

44 {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}

45 )

42 -@triton.jit

46 +@triton.jit(do_not_specialize=["N"])

47 def _selective_scan_update_kernel(

48 # Pointers to matrices

49 state_ptr,

@@ -55,8 +59,10 @@ def _selective_scan_update_kernel(

59 state_batch_indices_ptr,

60 dst_state_batch_indices_ptr,

61 pad_slot_id,

62 + num_accepted_tokens_ptr,

63 + cu_seqlens_ptr,

64 # Matrix dimensions

59 - batch,

65 + N,

66 nheads,

67 dim,

68 dstate,

@@ -91,6 +97,10 @@ def _selective_scan_update_kernel(

97 stride_out_batch,

98 stride_out_head,

99 stride_out_dim,

100 + stride_state_indices_batch,

101 + stride_state_indices_T,

102 + stride_dst_state_indices_batch,

103 + stride_dst_state_indices_T,

104 # Meta-parameters

105 DT_SOFTPLUS: tl.constexpr,

106 TIE_HDIM: tl.constexpr,

@@ -99,22 +109,50 @@ def _selective_scan_update_kernel(

109 HAS_D: tl.constexpr,

110 HAS_Z: tl.constexpr,

111 HAS_STATE_BATCH_INDICES: tl.constexpr,

112 + IS_SPEC_DECODING: tl.constexpr,

113 + IS_VARLEN: tl.constexpr,

114 BLOCK_SIZE_DSTATE: tl.constexpr,

115 ):

116 pid_m = tl.program_id(axis=0)

117 pid_b = tl.program_id(axis=1)

118 pid_h = tl.program_id(axis=2)

119

120 + if IS_VARLEN:

121 + bos = tl.load(cu_seqlens_ptr + pid_b).to(tl.int64)

122 + eos = tl.load(cu_seqlens_ptr + pid_b + 1).to(tl.int64)

123 + seq_len = eos - bos

124 +

125 + if seq_len == 0:

126 + return

127 + else:

128 + bos = pid_b

129 + seq_len = 1

130 +

131 + state_ptr_base = state_ptr

132 +

133 # If HAS_STATE_BATCH_INDICES is true, then the ssm state's batch coordinate

134 # is taken from the state_batch_indices_ptr Otherwise, the state coordinate

135 # is the same as the batch id.

136 if HAS_STATE_BATCH_INDICES:

112 - dst_state_batch_indices_ptr += pid_b

113 - dst_state_batch_idx = tl.load(dst_state_batch_indices_ptr).to(tl.int64)

114 - dst_state_ptr = state_ptr + (

115 - dst_state_batch_idx * stride_state_batch + pid_h * stride_state_head

137 + if IS_SPEC_DECODING:

138 + num_accepted = tl.load(num_accepted_tokens_ptr + pid_b).to(tl.int64)

139 + init_token_idx = tl.maximum(num_accepted - 1, 0)

140 + else:

141 + init_token_idx = 0

142 +

143 + dst_state_batch_indices_ptr += pid_b * stride_dst_state_indices_batch

144 + if not IS_SPEC_DECODING:

145 + dst_state_batch_idx = tl.load(

146 + dst_state_batch_indices_ptr

147 + + init_token_idx * stride_dst_state_indices_T

148 + ).to(tl.int64)

149 + dst_state_ptr = state_ptr + (

150 + dst_state_batch_idx * stride_state_batch + pid_h * stride_state_head

151 + )

152 +

153 + state_batch_indices_ptr += (

154 + pid_b * stride_state_indices_batch + init_token_idx * stride_state_indices_T

155 )

117 - state_batch_indices_ptr += pid_b

156 state_batch_idx = tl.load(state_batch_indices_ptr).to(tl.int64)

157 state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head

158 else:

@@ -123,86 +161,112 @@ def _selective_scan_update_kernel(

161 )

162 state_ptr += pid_b * stride_state_batch + pid_h * stride_state_head

163

126 - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head

127 - dt_ptr += pid_b * stride_dt_batch + pid_h * stride_dt_head

164 + x_ptr += bos * stride_x_batch + pid_h * stride_x_head

165 + dt_ptr += bos * stride_dt_batch + pid_h * stride_dt_head

166 if HAS_DT_BIAS:

167 dt_bias_ptr += pid_h * stride_dt_bias_head

168 A_ptr += pid_h * stride_A_head

131 - B_ptr += pid_b * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group

132 - C_ptr += pid_b * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group

169 + B_ptr += bos * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group

170 + C_ptr += bos * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group

171 if HAS_Z:

134 - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head

135 - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head

172 + z_ptr += bos * stride_z_batch + pid_h * stride_z_head

173 + out_ptr += bos * stride_out_batch + pid_h * stride_out_head

174

175 offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)

176 offs_n = tl.arange(0, BLOCK_SIZE_DSTATE)

177 state_ptrs = state_ptr + (

178 offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate

179 )

142 - dst_state_ptrs = dst_state_ptr + (

143 - offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate

144 - )

145 - x_ptrs = x_ptr + offs_m * stride_x_dim

146 - dt_ptrs = dt_ptr + offs_m * stride_dt_dim

180 + if not IS_SPEC_DECODING:

181 + dst_state_ptrs = dst_state_ptr + (

182 + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate

183 + )

184 +

185 + mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate)

186 + if HAS_STATE_BATCH_INDICES:

187 + mask &= state_batch_idx != pad_slot_id

188 + state = tl.load(state_ptrs, mask=mask, other=0.0).to(tl.float32)

189 +

190 if HAS_DT_BIAS:

191 dt_bias_ptrs = dt_bias_ptr + offs_m * stride_dt_bias_dim

192 if HAS_D:

193 D_ptr += pid_h * stride_D_head

151 - A_ptrs = A_ptr + (

152 - offs_m[:, None] * stride_A_dim + offs_n[None, :] * stride_A_dstate

153 - )

154 - B_ptrs = B_ptr + offs_n * stride_B_dstate

155 - C_ptrs = C_ptr + offs_n * stride_C_dstate

156 - if HAS_D:

194 D_ptrs = D_ptr + offs_m * stride_D_dim

158 - if HAS_Z:

159 - z_ptrs = z_ptr + offs_m * stride_z_dim

160 - out_ptrs = out_ptr + offs_m * stride_out_dim

161 - mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate)

162 - if HAS_STATE_BATCH_INDICES:

163 - mask &= state_batch_idx != pad_slot_id

164 - state = tl.load(state_ptrs, mask=mask, other=0.0)

165 -

166 - x = tl.load(x_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

167 - if not TIE_HDIM:

168 - dt = tl.load(dt_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

169 - if HAS_DT_BIAS:

170 - dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

171 - if DT_SOFTPLUS:

172 - dt = softplus(dt)

173 - A = tl.load(

174 - A_ptrs, mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), other=0.0

175 - ).to(tl.float32)

176 - dA = tl.exp(A * dt[:, None])

177 - else:

178 - dt = tl.load(dt_ptr).to(tl.float32)

179 - if HAS_DT_BIAS:

180 - dt += tl.load(dt_bias_ptr).to(tl.float32)

181 - if DT_SOFTPLUS:

182 - dt = softplus(dt)

183 - A = tl.load(A_ptr).to(tl.float32)

184 - dA = tl.exp(A * dt) # scalar, not a matrix

185 -

186 - B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)

187 - C = tl.load(C_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)

188 - if HAS_D:

189 - D = tl.load(D_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

190 - if HAS_Z:

191 - z = tl.load(z_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

195 + A_ptrs = A_ptr + offs_m[:, None] * stride_A_dim + offs_n[None, :] * stride_A_dstate

196

193 - dB = B[None, :] * dt[:, None] if not TIE_HDIM else B * dt

194 - state = state * dA + dB * x[:, None]

197 + for i_t in range(seq_len):

198 + x_ptrs = x_ptr + offs_m * stride_x_dim

199 + dt_ptrs = dt_ptr + offs_m * stride_dt_dim

200 + B_ptrs = B_ptr + offs_n * stride_B_dstate

201 + C_ptrs = C_ptr + offs_n * stride_C_dstate

202 + if HAS_Z:

203 + z_ptrs = z_ptr + offs_m * stride_z_dim

204 + out_ptrs = out_ptr + offs_m * stride_out_dim

205

196 - mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate)

197 - if HAS_STATE_BATCH_INDICES:

198 - mask &= state_batch_idx != pad_slot_id

199 - tl.store(dst_state_ptrs, state, mask=mask)

200 - out = tl.sum(state * C[None, :], axis=1)

201 - if HAS_D:

202 - out += x * D

203 - if HAS_Z:

204 - out *= z * tl.sigmoid(z)

205 - tl.store(out_ptrs, out, mask=offs_m < dim)

206 + x = tl.load(x_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

207 + if not TIE_HDIM:

208 + dt = tl.load(dt_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

209 + if HAS_DT_BIAS:

210 + dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

211 + if DT_SOFTPLUS:

212 + dt = softplus(dt)

213 + A = tl.load(

214 + A_ptrs,

215 + mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate),

216 + other=0.0,

217 + ).to(tl.float32)

218 + dA = tl.exp(A * dt[:, None])

219 + else:

220 + dt = tl.load(dt_ptr).to(tl.float32)

221 + if HAS_DT_BIAS:

222 + dt += tl.load(dt_bias_ptr).to(tl.float32)

223 + if DT_SOFTPLUS:

224 + dt = softplus(dt)

225 + A = tl.load(A_ptr).to(tl.float32)

226 + dA = tl.exp(A * dt) # scalar, not a matrix

227 +

228 + B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)

229 + C = tl.load(C_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)

230 + if HAS_D:

231 + D = tl.load(D_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

232 + if HAS_Z:

233 + z = tl.load(z_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)

234 +

235 + dB = B[None, :] * dt[:, None] if not TIE_HDIM else B * dt

236 + state = state * dA + dB * x[:, None]

237 +

238 + if IS_SPEC_DECODING:

239 + dst_idx_ptr = dst_state_batch_indices_ptr + i_t * stride_dst_state_indices_T

240 + token_dst_idx = tl.load(dst_idx_ptr).to(tl.int64)

241 + if token_dst_idx != pad_slot_id:

242 + token_dst_ptrs = (

243 + state_ptr_base

244 + + token_dst_idx * stride_state_batch

245 + + pid_h * stride_state_head

246 + + offs_m[:, None] * stride_state_dim

247 + + offs_n[None, :] * stride_state_dstate

248 + )

249 + tl.store(

250 + token_dst_ptrs, state.to(token_dst_ptrs.dtype.element_ty), mask=mask

251 + )

252 +

253 + out = tl.sum(state * C[None, :], axis=1)

254 + if HAS_D:

255 + out += x * D

256 + if HAS_Z:

257 + out *= z * tl.sigmoid(z)

258 + tl.store(out_ptrs, out, mask=offs_m < dim)

259 +

260 + x_ptr += stride_x_batch

261 + dt_ptr += stride_dt_batch

262 + B_ptr += stride_B_batch

263 + C_ptr += stride_C_batch

264 + out_ptr += stride_out_batch

265 + if HAS_Z:

266 + z_ptr += stride_z_batch

267 +

268 + if not IS_SPEC_DECODING:

269 + tl.store(dst_state_ptrs, state.to(dst_state_ptrs.dtype.element_ty), mask=mask)

270

271

272 def selective_state_update(

@@ -220,6 +284,8 @@ def selective_state_update(

284 dst_state_batch_indices=None,

285 pad_slot_id=PAD_SLOT_ID,

286 out=None,

287 + num_accepted_tokens=None,

288 + cu_seqlens=None,

289 ):

290 """

291 Argument:

@@ -240,6 +306,11 @@ def selective_state_update(

306 indices 0 and 3

307 out: Preallocated ssm output tensor. Assume same shape as x.

308 In-place updated.

309 + num_accepted_tokens: (batch,)

310 + number of accepted tokens from previous verification step,

311 + tells the kernel which initial state to use

312 + cu_seqlens: (batch,)

313 + length per sequence, for variable length in speculative decoding cases

314 """

315 if state.dim() == 3:

316 state = state.unsqueeze(1)

@@ -261,9 +332,26 @@ def selective_state_update(

332 dt_bias = dt_bias.unsqueeze(0)

333 if out.dim() == 2:

334 out = out.unsqueeze(1)

335 + if num_accepted_tokens is not None:

336 + assert state_batch_indices is not None and state_batch_indices.dim() == 2

337 + assert dst_state_batch_indices is None or dst_state_batch_indices.dim() == 2

338 + if state_batch_indices is not None and state_batch_indices.dim() == 1:

339 + state_batch_indices = state_batch_indices.unsqueeze(1)

340 + if dst_state_batch_indices is not None and dst_state_batch_indices.dim() == 1:

341 + dst_state_batch_indices = dst_state_batch_indices.unsqueeze(1)

342

343 _, nheads, dim, dstate = state.shape

344 batch = x.shape[0]

345 + if cu_seqlens is not None:

346 + N = len(cu_seqlens) - 1

347 + # Only used to verify the shape of

348 + # state_batch_indices and dst_state_batch_indices

349 + max_seqlen = (

350 + state_batch_indices.size(-1) if state_batch_indices is not None else 1

351 + )

352 + else:

353 + N = batch

354 + max_seqlen = 1

355

356 assert x.shape == (batch, nheads, dim)

357 assert dt.shape == x.shape

@@ -279,16 +367,30 @@ def selective_state_update(

367 if dt_bias is not None:

368 assert dt_bias.shape == (nheads, dim)

369 if state_batch_indices is not None:

282 - assert state_batch_indices.shape == (batch,)

370 + assert state_batch_indices.shape[0] >= N

371 + assert state_batch_indices.shape[1] >= max_seqlen

372 if dst_state_batch_indices is not None:

284 - assert dst_state_batch_indices.shape == (batch,)

373 + assert dst_state_batch_indices.shape[0] >= N

374 + assert dst_state_batch_indices.shape[1] >= max_seqlen

375 else:

376 # revert to the default behavior of in-place state updates

377 dst_state_batch_indices = state_batch_indices

378 assert out.shape == x.shape

379 + if num_accepted_tokens is not None:

380 + assert num_accepted_tokens.shape == (N,)

381

290 - grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads)

382 + grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), N, nheads)

383 z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0)

384 + state_batch_indices_strides = (

385 + (state_batch_indices.stride(0), state_batch_indices.stride(1))

386 + if state_batch_indices is not None

387 + else (0, 0)

388 + )

389 + dst_state_batch_indices_strides = (

390 + (dst_state_batch_indices.stride(0), dst_state_batch_indices.stride(1))

391 + if dst_state_batch_indices is not None

392 + else (0, 0)

393 + )

394 # We don't want autotune since it will overwrite the state

395 # We instead tune by hand.

396 BLOCK_SIZE_M, num_warps = (

@@ -321,7 +423,9 @@ def selective_state_update(

423 state_batch_indices,

424 dst_state_batch_indices,

425 pad_slot_id,

324 - batch,

426 + num_accepted_tokens,

427 + cu_seqlens,

428 + N,

429 nheads,

430 dim,

431 dstate,

@@ -353,6 +457,10 @@ def selective_state_update(

457 out.stride(0),

458 out.stride(1),

459 out.stride(2),

460 + state_batch_indices_strides[0],

461 + state_batch_indices_strides[1],

462 + dst_state_batch_indices_strides[0],

463 + dst_state_batch_indices_strides[1],

464 dt_softplus,

465 tie_hdim,

466 BLOCK_SIZE_M,