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

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 5, 14:45 UTC

[GDN] Eliminate GPU->CPU sync in prepare_chunk_indices during prefill (#38361)

Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com> Signed-off-by: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com>

Artem Perevedentsev committed Apr 3, 2026 at 16:38 UTC cb10b7e80b138f876079c2cd697df29f0965b885

10 files changed +116 -44

vllm/model_executor/layers/fla/ops/chunk.py

+27 -4

@@ -16,7 +16,7 @@ from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd

16 from .cumsum import chunk_local_cumsum

17 from .l2norm import l2norm_fwd

18 from .solve_tril import solve_tril

19 -from .utils import SUPPRESS_LEVEL, input_guard

19 +from .utils import FLA_CHUNK_SIZE, SUPPRESS_LEVEL, input_guard

20 from .wy_fast import recompute_w_u_fwd

@@ -30,13 +30,24 @@ def chunk_gated_delta_rule_fwd(

30 initial_state: torch.Tensor,

31 output_final_state: bool,

32 cu_seqlens: torch.Tensor | None = None,

33 + chunk_indices: torch.Tensor | None = None,

34 + chunk_offsets: torch.Tensor | None = None,

35 ):

34 - g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens)

36 + g = chunk_local_cumsum(

37 + g, chunk_size=FLA_CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices

38 + )

39 # obtain WY representation. u is actually the new v.

40 A = chunk_scaled_dot_kkt_fwd(

37 - k=k, beta=beta, g=g, cu_seqlens=cu_seqlens, output_dtype=torch.float32

41 + k=k,

42 + beta=beta,

43 + g=g,

44 + cu_seqlens=cu_seqlens,

45 + chunk_indices=chunk_indices,

46 + output_dtype=torch.float32,

47 + )

48 + A = solve_tril(

49 + A=A, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, output_dtype=k.dtype

50 )

39 - A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)

51 w, u = recompute_w_u_fwd(

52 k=k,

53 v=v,

@@ -44,6 +55,7 @@ def chunk_gated_delta_rule_fwd(

55 A=A,

56 g_cumsum=g,

57 cu_seqlens=cu_seqlens,

58 + chunk_indices=chunk_indices,

59 )

60 h, v_new, final_state = chunk_gated_delta_rule_fwd_h(

61 k=k,

@@ -53,6 +65,8 @@ def chunk_gated_delta_rule_fwd(

65 initial_state=initial_state,

66 output_final_state=output_final_state,

67 cu_seqlens=cu_seqlens,

68 + chunk_indices=chunk_indices,

69 + chunk_offsets=chunk_offsets,

70 )

71 o = chunk_fwd_o(

72 q=q,

@@ -62,6 +76,7 @@ def chunk_gated_delta_rule_fwd(

76 g=g,

77 scale=scale,

78 cu_seqlens=cu_seqlens,

79 + chunk_indices=chunk_indices,

80 )

81 if SUPPRESS_LEVEL < 3:

82 return g, o, A, final_state, None, None, None

@@ -84,6 +99,8 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):

99 initial_state: torch.Tensor,

100 output_final_state: bool,

101 cu_seqlens: torch.Tensor | None = None,

102 + chunk_indices: torch.Tensor | None = None,

103 + chunk_offsets: torch.Tensor | None = None,

104 use_qk_l2norm_in_kernel: bool = False,

105 ):

106 if use_qk_l2norm_in_kernel:

@@ -100,6 +117,8 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):

117 initial_state=initial_state,

118 output_final_state=output_final_state,

119 cu_seqlens=cu_seqlens,

120 + chunk_indices=chunk_indices,

121 + chunk_offsets=chunk_offsets,

122 )

123 ctx.scale = scale

124 ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel

@@ -117,6 +136,8 @@ def chunk_gated_delta_rule(

136 initial_state: torch.Tensor = None,

137 output_final_state: bool = False,

138 cu_seqlens: torch.Tensor | None = None,

139 + chunk_indices: torch.Tensor | None = None,

140 + chunk_offsets: torch.Tensor | None = None,

141 use_qk_l2norm_in_kernel: bool = False,

142 ):

143 r"""

@@ -206,6 +227,8 @@ def chunk_gated_delta_rule(

227 initial_state,

228 output_final_state,

229 cu_seqlens,

230 + chunk_indices,

231 + chunk_offsets,

232 use_qk_l2norm_in_kernel,

233 )

234 return o, final_state

vllm/model_executor/layers/fla/ops/chunk_delta_h.py

+9 -12

@@ -14,7 +14,7 @@ from vllm.triton_utils import tl, triton

15 from .index import prepare_chunk_indices, prepare_chunk_offsets

16 from .op import exp

17 -from .utils import use_cuda_graph

17 +from .utils import FLA_CHUNK_SIZE, use_cuda_graph

19 NUM_WARPS = [2, 4, 8, 16]

@@ -286,9 +286,11 @@ def chunk_gated_delta_rule_fwd_h(

286 gk: torch.Tensor | None = None,

287 initial_state: torch.Tensor | None = None,

288 output_final_state: bool = False,

289 - chunk_size: int = 64, # SY: remove this argument and force chunk size 64?

289 + chunk_size: int = FLA_CHUNK_SIZE,

290 save_new_value: bool = True,

291 cu_seqlens: torch.Tensor | None = None,

292 + chunk_indices: torch.Tensor | None = None,

293 + chunk_offsets: torch.Tensor | None = None,

294 ) -> tuple[torch.Tensor, torch.Tensor]:

295 # This kernel is slightly different from fla to support Q/K with different head numbers.

296 # In fla, Q/K always have the same head number, so Hg is always equal to H.

@@ -296,20 +298,15 @@ def chunk_gated_delta_rule_fwd_h(

298 H = u.shape[-2]

299 BT = chunk_size

300

299 - chunk_indices = (

300 - prepare_chunk_indices(cu_seqlens, chunk_size)

301 - if cu_seqlens is not None

302 - else None

303 - )

301 + if chunk_indices is None and cu_seqlens is not None:

302 + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)

303 # N: the actual number of sequences in the batch with either equal or variable lengths

304 if cu_seqlens is None:

305 N, NT, chunk_offsets = B, triton.cdiv(T, BT), None

306 else:

308 - N, NT, chunk_offsets = (

309 - len(cu_seqlens) - 1,

310 - len(chunk_indices),

311 - prepare_chunk_offsets(cu_seqlens, BT),

312 - )

307 + N, NT = len(cu_seqlens) - 1, len(chunk_indices)

308 + if chunk_offsets is None:

309 + chunk_offsets = prepare_chunk_offsets(cu_seqlens, BT)

310 assert K <= 256, "current kernel does not support head dimension larger than 256."

311

312 h = k.new_empty(B, NT, H, V, K)

vllm/model_executor/layers/fla/ops/chunk_o.py

+3 -3

@@ -146,14 +146,14 @@ def chunk_fwd_o(

146 g: torch.Tensor | None = None, # cumsum of log decay

147 scale: float | None = None,

148 cu_seqlens: torch.Tensor | None = None,

149 + chunk_indices: torch.Tensor | None = None,

150 chunk_size: int = FLA_CHUNK_SIZE,

151 ) -> torch.Tensor:

152 B, T, Hg, K, V = *q.shape, v.shape[-1]

153 H = v.shape[-2]

154 BT = chunk_size

154 - chunk_indices = (

155 - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None

156 - )

155 + if chunk_indices is None and cu_seqlens is not None:

156 + chunk_indices = prepare_chunk_indices(cu_seqlens, BT)

157 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

158 if scale is None:

159 scale = k.shape[-1] ** -0.5

vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py

+8 -4

@@ -14,6 +14,7 @@ from vllm.triton_utils import tl, triton

15 from .index import prepare_chunk_indices

16 from .op import exp

17 +from .utils import FLA_CHUNK_SIZE

20 @triton.heuristics(

@@ -103,7 +104,8 @@ def chunk_scaled_dot_kkt_fwd(

104 g: torch.Tensor | None = None,

105 beta: torch.Tensor | None = None,

106 cu_seqlens: torch.Tensor | None = None,

106 - chunk_size: int = 64,

107 + chunk_indices: torch.Tensor | None = None,

108 + chunk_size: int = FLA_CHUNK_SIZE,

109 output_dtype: torch.dtype = torch.float32,

110 ) -> torch.Tensor:

111 r"""

@@ -119,6 +121,9 @@ def chunk_scaled_dot_kkt_fwd(

121 cu_seqlens (torch.Tensor):

122 The cumulative sequence lengths of the input tensor.

123 Default: None

124 + chunk_indices (torch.Tensor):

125 + Pre-computed chunk indices. If None and cu_seqlens is provided,

126 + computed internally. Default: None

127 chunk_size (int):

128 The chunk size. Default: 64.

129 output_dtype (torch.dtype):

@@ -132,9 +137,8 @@ def chunk_scaled_dot_kkt_fwd(

137 B, T, Hg, K = k.shape

138 H = beta.shape[-1]

139 BT = chunk_size

135 - chunk_indices = (

136 - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None

137 - )

140 + if chunk_indices is None and cu_seqlens is not None:

141 + chunk_indices = prepare_chunk_indices(cu_seqlens, BT)

142 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

143

144 A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)

vllm/model_executor/layers/fla/ops/cumsum.py

+23 -12

@@ -162,6 +162,7 @@ def chunk_local_cumsum_scalar(

162 chunk_size: int,

163 reverse: bool = False,

164 cu_seqlens: torch.Tensor | None = None,

165 + chunk_indices: torch.Tensor | None = None,

166 head_first: bool = False,

167 output_dtype: torch.dtype | None = torch.float,

168 ) -> torch.Tensor:

@@ -172,10 +173,9 @@ def chunk_local_cumsum_scalar(

173 assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (

174 "chunk_size must be a power of 2"

175 )

176 + if chunk_indices is None and cu_seqlens is not None:

177 + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)

178 BT = chunk_size

176 - chunk_indices = (

177 - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None

178 - )

179 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

180 g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)

181 grid = (NT, B * H)

@@ -199,6 +199,7 @@ def chunk_local_cumsum_vector(

199 chunk_size: int,

200 reverse: bool = False,

201 cu_seqlens: torch.Tensor | None = None,

202 + chunk_indices: torch.Tensor | None = None,

203 head_first: bool = False,

204 output_dtype: torch.dtype | None = torch.float,

205 ) -> torch.Tensor:

@@ -206,16 +207,13 @@ def chunk_local_cumsum_vector(

207 B, H, T, S = g.shape

208 else:

209 B, T, H, S = g.shape

209 - BT = chunk_size

210 - chunk_indices = (

211 - prepare_chunk_indices(cu_seqlens, chunk_size)

212 - if cu_seqlens is not None

213 - else None

214 - )

215 - NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

210 assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (

211 "chunk_size must be a power of 2"

212 )

213 + if chunk_indices is None and cu_seqlens is not None:

214 + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)

215 + BT = chunk_size

216 + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

217

218 g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)

219

@@ -247,6 +245,7 @@ def chunk_local_cumsum(

245 chunk_size: int,

246 reverse: bool = False,

247 cu_seqlens: torch.Tensor | None = None,

248 + chunk_indices: torch.Tensor | None = None,

249 head_first: bool = False,

250 output_dtype: torch.dtype | None = torch.float,

251 **kwargs,

@@ -257,11 +256,23 @@ def chunk_local_cumsum(

256 )

257 if len(g.shape) == 3:

258 return chunk_local_cumsum_scalar(

260 - g, chunk_size, reverse, cu_seqlens, head_first, output_dtype

259 + g,

260 + chunk_size,

261 + reverse,

262 + cu_seqlens,

263 + chunk_indices,

264 + head_first,

265 + output_dtype,

266 )

267 elif len(g.shape) == 4:

268 return chunk_local_cumsum_vector(

264 - g, chunk_size, reverse, cu_seqlens, head_first, output_dtype

269 + g,

270 + chunk_size,

271 + reverse,

272 + cu_seqlens,

273 + chunk_indices,

274 + head_first,

275 + output_dtype,

276 )

277 else:

278 raise ValueError(

vllm/model_executor/layers/fla/ops/kda.py

+4 -3

@@ -23,7 +23,7 @@ from .index import prepare_chunk_indices

23 from .l2norm import l2norm_fwd

24 from .op import exp, log

25 from .solve_tril import solve_tril

26 -from .utils import is_amd

26 +from .utils import FLA_CHUNK_SIZE, is_amd

28 BT_LIST_AUTOTUNE = [32, 64, 128]

29 NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]

@@ -721,7 +721,7 @@ def chunk_kda_scaled_dot_kkt_fwd(

721 beta: torch.Tensor | None = None,

722 scale: float | None = None,

723 cu_seqlens: torch.Tensor | None = None,

724 - chunk_size: int = 64,

724 + chunk_size: int = FLA_CHUNK_SIZE,

725 output_dtype: torch.dtype = torch.float32,

726 ) -> tuple[torch.Tensor, torch.Tensor]:

727 r"""

@@ -1178,7 +1178,7 @@ def chunk_kda_fwd(

1178 output_final_state: bool,

1179 cu_seqlens: torch.Tensor | None = None,

1180 ):

1181 - chunk_size = 64

1181 + chunk_size = FLA_CHUNK_SIZE

1182 g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens)

1183 # the intra Aqk is kept in fp32

1184 # the computation has very marginal effect on the entire throughput

@@ -1189,6 +1189,7 @@ def chunk_kda_fwd(

1189 beta=beta,

1190 scale=scale,

1191 cu_seqlens=cu_seqlens,

1192 + chunk_size=chunk_size,

1193 output_dtype=torch.float32,

1194 )

1195 A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)

vllm/model_executor/layers/fla/ops/solve_tril.py

+5 -3

@@ -507,6 +507,7 @@ def merge_16x16_to_64x64_inverse_kernel(

507 def solve_tril(

508 A: torch.Tensor,

509 cu_seqlens: torch.Tensor | None = None,

510 + chunk_indices: torch.Tensor | None = None,

511 output_dtype: torch.dtype = torch.float,

512 ) -> torch.Tensor:

513 """

@@ -518,6 +519,8 @@ def solve_tril(

519 [B, T, H, BT], where BT should only be 16, 32, or 64.

520 cu_seqlens (torch.Tensor):

521 The cumulative sequence lengths of the input tensor. Default: `None`.

522 + chunk_indices (torch.Tensor):

523 + Pre-computed chunk indices. Default: `None`.

524 output_dtype (torch.dtype):

525 The dtype of the output tensor. Default: `torch.float`.

526 If `None`, the output dtype will be the same as the input dtype.

@@ -529,9 +532,8 @@ def solve_tril(

532 output_dtype = A.dtype if output_dtype is None else output_dtype

533

534 B, T, H, BT = A.shape

532 - chunk_indices = (

533 - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None

534 - )

535 + if chunk_indices is None and cu_seqlens is not None:

536 + chunk_indices = prepare_chunk_indices(cu_seqlens, BT)

537 NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)

538

539 Ai = torch.zeros_like(A, dtype=output_dtype)

vllm/model_executor/layers/fla/ops/wy_fast.py

+3 -3

@@ -123,14 +123,14 @@ def recompute_w_u_fwd(

123 g_cumsum: torch.Tensor,

124 A: torch.Tensor,

125 cu_seqlens: torch.Tensor | None,

126 + chunk_indices: torch.Tensor | None = None,

127 ) -> tuple[torch.Tensor, torch.Tensor]:

128 B, T, Hg, K, V = *k.shape, v.shape[-1]

129 H = v.shape[-2]

130 BT = A.shape[-1]

131

131 - chunk_indices = (

132 - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None

133 - )

132 + if chunk_indices is None and cu_seqlens is not None:

133 + chunk_indices = prepare_chunk_indices(cu_seqlens, BT)

134 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

135 BK = 64

136 BV = 64

vllm/model_executor/layers/mamba/gdn_linear_attn.py

@@ -163,6 +163,8 @@ class ChunkGatedDeltaRule(CustomOp):

163 initial_state: torch.Tensor,

164 output_final_state: bool,

165 cu_seqlens: torch.Tensor | None = None,

166 + chunk_indices: torch.Tensor | None = None,

167 + chunk_offsets: torch.Tensor | None = None,

168 use_qk_l2norm_in_kernel: bool = True,

169 ):

170 return fi_chunk_gated_delta_rule(

@@ -187,6 +189,8 @@ class ChunkGatedDeltaRule(CustomOp):

189 initial_state: torch.Tensor,

190 output_final_state: bool,

191 cu_seqlens: torch.Tensor | None = None,

192 + chunk_indices: torch.Tensor | None = None,

193 + chunk_offsets: torch.Tensor | None = None,

194 use_qk_l2norm_in_kernel: bool = True,

195 ):

196 return fla_chunk_gated_delta_rule(

@@ -198,6 +202,8 @@ class ChunkGatedDeltaRule(CustomOp):

202 initial_state=initial_state,

203 output_final_state=output_final_state,

204 cu_seqlens=cu_seqlens,

205 + chunk_indices=chunk_indices,

206 + chunk_offsets=chunk_offsets,

207 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,

208 )

209

@@ -959,6 +965,8 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):

965 initial_state=initial_state,

966 output_final_state=True,

967 cu_seqlens=non_spec_query_start_loc,

968 + chunk_indices=attn_metadata.chunk_indices,

969 + chunk_offsets=attn_metadata.chunk_offsets,

970 use_qk_l2norm_in_kernel=False,

971 )

972 # Init cache

vllm/v1/attention/backends/gdn_attn.py

+26

@@ -63,6 +63,10 @@ class GDNAttentionMetadata:

64 num_accepted_tokens: torch.Tensor | None = None # shape: [batch,]

66 + # Pre-computed FLA chunk metadata (avoids GPU->CPU sync in prepare_chunk_indices)

67 + chunk_indices: torch.Tensor | None = None

68 + chunk_offsets: torch.Tensor | None = None

69 +

70 # The following attributes are for triton implementation of causal_conv1d

71 nums_dict: dict | None = None

72 batch_ptr: torch.Tensor | None = None

@@ -305,6 +309,26 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

309 assert num_accepted_tokens is not None

310 num_accepted_tokens = num_accepted_tokens[spec_sequence_masks]

311

312 + chunk_indices: torch.Tensor | None = None

313 + chunk_offsets: torch.Tensor | None = None

314 + if num_prefills > 0:

315 + # Only prefill batches use FLA chunk ops.

316 + # Pre-compute on CPU and async-copy to GPU to avoid

317 + # GPU→CPU sync (.tolist()) in prepare_chunk_indices.

318 + from vllm.model_executor.layers.fla.ops.index import (

319 + prepare_chunk_indices,

320 + prepare_chunk_offsets,

321 + )

322 + from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE

323 +

324 + gpu_device = query_start_loc.device

325 + chunk_indices = prepare_chunk_indices(

326 + non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE

327 + ).to(device=gpu_device, non_blocking=True)

328 + chunk_offsets = prepare_chunk_offsets(

329 + non_spec_query_start_loc_cpu, FLA_CHUNK_SIZE

330 + ).to(device=gpu_device, non_blocking=True)

331 +

332 if num_prefills > 0:

333 has_initial_state = context_lens_tensor > 0

334 if spec_sequence_masks is not None:

@@ -405,6 +429,8 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

429 num_spec_decode_tokens=num_spec_decode_tokens,

430 num_actual_tokens=m.num_actual_tokens,

431 has_initial_state=has_initial_state,

432 + chunk_indices=chunk_indices,

433 + chunk_offsets=chunk_offsets,

434 spec_query_start_loc=spec_query_start_loc,

435 non_spec_query_start_loc=non_spec_query_start_loc,

436 spec_state_indices_tensor=spec_state_indices_tensor,