[Refactor] Remove unused rocm kernel `combine_topk_swa_indices_ragged` (#48158)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 10, 2026 at 09:33 UTC
e257faf87d8e003451e8693262c2f90df4acea45
2 files changed
-236
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
-89
@@ -197,46 +197,6 @@ def _ragged_from_rows(
197
)
198
199
200
-def _ref_combine_topk_swa_ragged(
201
- device: torch.device,
202
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
203
- expected_ragged = torch.tensor(
204
- [
205
- 100,
206
- 101,
207
- 7,
208
- 8,
209
- 9,
210
- 110,
211
- 111,
212
- 8,
213
- 9,
214
- 10,
215
- 120,
216
- 121,
217
- 122,
218
- 9,
219
- 10,
220
- 11,
221
- 150,
222
- 27,
223
- 28,
224
- 29,
225
- 160,
226
- 161,
227
- 28,
228
- 29,
229
- 30,
230
- ],
231
- dtype=torch.int32,
232
- device=device,
233
- )
234
- expected_lens = torch.tensor([5, 5, 6, 4, 5], dtype=torch.int32, device=device)
235
- expected_indptr = torch.zeros(6, dtype=torch.int32, device=device)
236
- torch.cumsum(expected_lens, dim=0, out=expected_indptr[1:])
237
- return expected_ragged, expected_indptr, expected_lens
238
-
239
-
200
@torch.inference_mode()
201
def test_compute_global_topk_ragged_indices_and_indptr() -> None:
202
from vllm.models.deepseek_v4.amd.rocm import (
@@ -369,55 +329,6 @@ def test_sparse_attn_decode_ragged_kernel() -> None:
329
torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
330
331
372
-@torch.inference_mode()
373
-def test_combine_topk_swa_indices_ragged() -> None:
374
- from vllm.models.deepseek_v4.amd.rocm import (
375
- combine_topk_swa_indices_ragged,
376
- )
377
-
378
- device = torch.device("cuda")
379
- topk_indices = torch.tensor(
380
- [
381
- [100, 101, 102, 103],
382
- [110, 111, 112, 113],
383
- [120, 121, 122, 123],
384
- [130, 131, 132, 133],
385
- [140, 141, 142, 143],
386
- ],
387
- dtype=torch.int32,
388
- device=device,
389
- )
390
- query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
391
- seq_lens = torch.tensor([6, 4], dtype=torch.int32, device=device)
392
- gather_lens = torch.tensor([4, 3], dtype=torch.int32, device=device)
393
- window_size = 3
394
- compress_ratio = 2
395
- topk = 4
396
- M = 20
397
- N = 8
398
-
399
- actual_ragged, actual_indptr, actual_lens = combine_topk_swa_indices_ragged(
400
- topk_indices,
401
- query_start_loc,
402
- seq_lens,
403
- gather_lens,
404
- window_size,
405
- compress_ratio,
406
- topk,
407
- M,
408
- N,
409
- )
410
- expected_ragged, expected_indptr, expected_lens = _ref_combine_topk_swa_ragged(
411
- device
412
- )
413
-
414
- torch.testing.assert_close(
415
- actual_ragged[: expected_ragged.numel()], expected_ragged
416
- )
417
- torch.testing.assert_close(actual_indptr, expected_indptr)
418
- torch.testing.assert_close(actual_lens, expected_lens)
419
-
420
-
332
@requires_gfx950
333
@torch.inference_mode()
334
def test_decode_num_splits_heuristic(monkeypatch) -> None:
vllm/models/deepseek_v4/amd/rocm.py
-147
@@ -272,153 +272,6 @@ def compute_global_topk_ragged_indices_and_indptr(
272
return global_topk_ragged, topk_indptr, topk_lens
273
274
275
-@triton.jit
276
-def _compute_combined_lens_kernel(
277
- combined_lens_ptr,
278
- query_start_loc_ptr,
279
- seq_lens_ptr,
280
- TOP_K: tl.constexpr,
281
- COMPRESS_RATIO: tl.constexpr,
282
- WINDOW_SIZE: tl.constexpr,
283
-):
284
- batch_idx = tl.program_id(0)
285
- worker_id = tl.program_id(1)
286
- num_workers = tl.num_programs(1)
287
-
288
- base = tl.load(query_start_loc_ptr)
289
- query_start = tl.load(query_start_loc_ptr + batch_idx) - base
290
- query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
291
- query_len = query_end - query_start
292
- seq_len = tl.load(seq_lens_ptr + batch_idx)
293
- start_pos = seq_len - query_len
294
-
295
- for token_idx in range(query_start + worker_id, query_end, num_workers):
296
- token_idx_in_query = token_idx - query_start
297
- pos = start_pos + token_idx_in_query
298
- topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
299
- swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
300
- tl.store(combined_lens_ptr + token_idx, topk_len + swa_len)
301
-
302
-
303
-@triton.jit
304
-def _combine_topk_swa_indices_ragged_kernel(
305
- combined_ragged_ptr,
306
- combined_indptr_ptr,
307
- topk_indices_ptr,
308
- topk_indices_stride,
309
- query_start_loc_ptr,
310
- seq_lens_ptr,
311
- gather_lens_ptr,
312
- M,
313
- N,
314
- topk_width,
315
- TOP_K: tl.constexpr,
316
- COMPRESS_RATIO: tl.constexpr,
317
- WINDOW_SIZE: tl.constexpr,
318
- BLOCK_SIZE: tl.constexpr,
319
-):
320
- batch_idx = tl.program_id(0)
321
- worker_id = tl.program_id(1)
322
- block_idx = tl.program_id(2)
323
- num_workers = tl.num_programs(1)
324
-
325
- base = tl.load(query_start_loc_ptr)
326
- query_start = tl.load(query_start_loc_ptr + batch_idx) - base
327
- query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
328
- query_len = query_end - query_start
329
- seq_len = tl.load(seq_lens_ptr + batch_idx)
330
- gather_len = tl.load(gather_lens_ptr + batch_idx)
331
- start_pos = seq_len - query_len
332
- gather_start = seq_len - gather_len
333
-
334
- for token_idx in range(query_start + worker_id, query_end, num_workers):
335
- token_idx_in_query = token_idx - query_start
336
- pos = start_pos + token_idx_in_query
337
- topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
338
- swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
339
- combined_len = topk_len + swa_len
340
-
341
- offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
342
- if block_idx * BLOCK_SIZE < combined_len:
343
- out_start = tl.load(combined_indptr_ptr + token_idx)
344
- topk_mask = (offset < topk_len) & (offset < topk_width)
345
- topk_vals = tl.load(
346
- topk_indices_ptr + token_idx * topk_indices_stride + offset,
347
- mask=topk_mask,
348
- other=-1,
349
- )
350
- tl.store(
351
- combined_ragged_ptr + out_start + offset,
352
- topk_vals + M * batch_idx,
353
- mask=topk_mask,
354
- )
355
-
356
- swa_offset = offset - topk_len
357
- swa_mask = (offset >= topk_len) & (swa_offset < swa_len)
358
- tl.store(
359
- combined_ragged_ptr + out_start + offset,
360
- M * batch_idx + N + swa_offset + pos - swa_len + 1 - gather_start,
361
- mask=swa_mask,
362
- )
363
-
364
-
365
-def combine_topk_swa_indices_ragged(
366
- topk_indices: torch.Tensor,
367
- query_start_loc: torch.Tensor,
368
- seq_lens: torch.Tensor,
369
- gather_lens: torch.Tensor,
370
- window_size: int,
371
- compress_ratio: int,
372
- topk: int,
373
- M: int,
374
- N: int,
375
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
376
- topk_indices = topk_indices.reshape(topk_indices.shape[0], -1).contiguous()
377
- num_tokens = topk_indices.shape[0]
378
- num_reqs = seq_lens.shape[0]
379
- combined_lens = torch.empty(
380
- num_tokens, dtype=torch.int32, device=topk_indices.device
381
- )
382
-
383
- num_workers = 128
384
- _compute_combined_lens_kernel[(num_reqs, num_workers)](
385
- combined_lens,
386
- query_start_loc,
387
- seq_lens,
388
- TOP_K=topk,
389
- COMPRESS_RATIO=compress_ratio,
390
- WINDOW_SIZE=window_size,
391
- )
392
-
393
- combined_indptr = _build_indptr_from_lengths(combined_lens)
394
- combined_ragged = torch.empty(
395
- num_tokens * (topk + window_size),
396
- dtype=torch.int32,
397
- device=topk_indices.device,
398
- )
399
- if combined_ragged.numel() > 0:
400
- block = 128
401
- _combine_topk_swa_indices_ragged_kernel[
402
- (num_reqs, num_workers, triton.cdiv(topk + window_size, block))
403
- ](
404
- combined_ragged,
405
- combined_indptr,
406
- topk_indices,
407
- topk_indices.stride(0),
408
- query_start_loc,
409
- seq_lens,
410
- gather_lens,
411
- M,
412
- N,
413
- topk_indices.shape[-1],
414
- TOP_K=topk,
415
- COMPRESS_RATIO=compress_ratio,
416
- WINDOW_SIZE=window_size,
417
- BLOCK_SIZE=block,
418
- )
419
- return combined_ragged, combined_indptr, combined_lens
420
-
421
-
275
def _copy_ragged_to_graph_buffers(
276
ragged_indices: torch.Tensor,
277
ragged_indptr: torch.Tensor,