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 15, 04:45 UTC
[Refactor] Remove duplicate `ceil_div` (#20023)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 25, 2025 at 01:19 UTC
879f69bed375968ce02e003883752754beb36111
7 files changed
+20
-42
benchmarks/cutlass_benchmarks/w8a8_benchmarks.py
+3
-8
@@ -19,7 +19,7 @@ from vllm import _custom_ops as ops
19
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
20
w8a8_block_fp8_matmul,
21
)
22
-from vllm.utils import FlexibleArgumentParser
22
+from vllm.utils import FlexibleArgumentParser, cdiv
23
24
DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
25
DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]
@@ -117,14 +117,9 @@ def bench_fp8(
117
scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
118
scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
119
120
- def ceil_div(x: int, y: int) -> int:
121
- return (x + y - 1) // y
122
-
123
- block_scale_a = torch.rand(
124
- (m, ceil_div(k, 128)), device="cuda", dtype=torch.float32
125
- )
120
+ block_scale_a = torch.rand((m, cdiv(k, 128)), device="cuda", dtype=torch.float32)
121
block_scale_b = torch.rand(
127
- ceil_div(k, 128), ceil_div(n, 128), device="cuda", dtype=torch.float32
122
+ cdiv(k, 128), cdiv(n, 128), device="cuda", dtype=torch.float32
123
)
124
block_scale_a_M_major = block_scale_a.t().contiguous().t()
125
block_scale_b_K_major = block_scale_b.t().contiguous().t()
tests/kernels/attention/test_mla_decode_cpu.py
+1
-4
@@ -7,10 +7,7 @@ from torch import Tensor
7
8
import vllm._custom_ops as ops
9
from vllm.platforms import current_platform
10
-
11
-
12
-def cdiv(a, b):
13
- return (a + b - 1) // b
10
+from vllm.utils import cdiv
11
12
13
def ref_mla(
tests/kernels/attention/test_triton_decode_attention.py
+1
-4
@@ -5,10 +5,7 @@ import pytest
5
import torch
6
7
from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
8
-
9
-
10
-def cdiv(a, b):
11
- return (a + b - 1) // b
8
+from vllm.utils import cdiv
9
10
11
@pytest.mark.parametrize("B", [3, 5])
tests/neuron/1_core/test_prefix_prefill.py
+4
-5
@@ -7,6 +7,8 @@ import pytest
7
import torch
8
import torch.nn.functional as F
9
10
+from vllm.utils import cdiv
11
+
12
13
class BlockDiagonalCausalFromBottomRightMask:
14
@@ -398,11 +400,8 @@ def test_contexted_kv_attention(
400
assert (large_tile_size >= B_P_SIZE
401
), f"Expect {large_tile_size=} to be larger than {B_P_SIZE=}"
402
401
- def ceil_div(a, b):
402
- return (a + b - 1) // b
403
-
403
def pad_to_multiple(a, b):
405
- return ceil_div(a, b) * b
404
+ return cdiv(a, b) * b
405
406
def pad_to_next_power_of_2(a):
407
assert a > 0
@@ -411,7 +410,7 @@ def test_contexted_kv_attention(
410
# calculate input shapes
411
max_num_queries = pad_to_next_power_of_2(sum(query_lens))
412
context_lens = torch.tensor(seq_lens) - torch.tensor(query_lens)
414
- num_active_blocks = ceil_div(context_lens, block_size).sum().item()
413
+ num_active_blocks = cdiv(context_lens, block_size).sum().item()
414
num_active_blocks = pad_to_multiple(num_active_blocks,
415
large_tile_size // block_size)
416
context_kv_len = num_active_blocks * block_size
vllm/attention/ops/nki_flash_attn.py
+6
-9
@@ -8,9 +8,7 @@ import torch
8
from neuronxcc import nki
9
from neuronxcc.nki.language import par_dim
10
11
-
12
-def ceil_div(a, b):
13
- return (a + b - 1) // b
11
+from vllm.utils import cdiv
12
13
14
def is_power_of_2(x):
@@ -35,11 +33,10 @@ def load_block_tables(block_tables_hbm, num_tiles, num_blocks_per_tile):
33
(num_tiles, num_blocks_per_tile))
34
35
block_tables_sbuf = nl.zeros(
38
- (ceil_div(num_tiles,
39
- B_P_SIZE), par_dim(B_P_SIZE), num_blocks_per_tile),
36
+ (cdiv(num_tiles, B_P_SIZE), par_dim(B_P_SIZE), num_blocks_per_tile),
37
dtype=nl.int32,
38
)
42
- for i in nl.affine_range(ceil_div(num_tiles, B_P_SIZE)):
39
+ for i in nl.affine_range(cdiv(num_tiles, B_P_SIZE)):
40
i_p = nl.arange(B_P_SIZE)[:, None]
41
i_f = nl.arange(num_blocks_per_tile)[None, :]
42
block_tables_sbuf[i, i_p, i_f] = nl.load(
@@ -83,7 +80,7 @@ def transform_block_tables_for_indirect_load(
80
assert is_power_of_2(
81
num_blocks_per_tile), f"{num_blocks_per_tile=} is not power of 2"
82
86
- num_loads = ceil_div(num_blocks_per_tile, B_P_SIZE)
83
+ num_loads = cdiv(num_blocks_per_tile, B_P_SIZE)
84
block_tables_transposed = nl.ndarray(
85
(
86
num_loads,
@@ -165,7 +162,7 @@ def load_kv_tile_from_cache(
162
equivalent to (par_dim(B_P_SIZE), seqlen_kv // B_P_SIZE * B_D_SIZE)
163
"""
164
# load key cache
168
- num_loads = ceil_div(num_blocks_per_large_tile, B_P_SIZE)
165
+ num_loads = cdiv(num_blocks_per_large_tile, B_P_SIZE)
166
for load_idx in nl.affine_range(num_loads):
167
i_p = nl.arange(B_P_SIZE)[:, None]
168
i_f = nl.arange(tiled_block_size * B_D_SIZE)[None, :]
@@ -605,7 +602,7 @@ def flash_paged_attention(
602
)
603
604
for large_k_tile_idx in nl.sequential_range(0, num_large_k_tile):
608
- num_loads = ceil_div(num_blocks_per_large_tile, B_P_SIZE)
605
+ num_loads = cdiv(num_blocks_per_large_tile, B_P_SIZE)
606
cur_k_tile = nl.ndarray(
607
(par_dim(B_D_SIZE), LARGE_TILE_SZ),
608
dtype=kernel_dtype,
vllm/model_executor/layers/fused_moe/moe_align_block_size.py
+2
-6
@@ -6,11 +6,7 @@ import torch
6
7
from vllm import _custom_ops as ops
8
from vllm.triton_utils import tl, triton
9
-from vllm.utils import round_up
10
-
11
-
12
-def ceil_div(a, b):
13
- return (a + b - 1) // b
9
+from vllm.utils import cdiv, round_up
10
11
12
@triton.jit
@@ -115,7 +111,7 @@ def moe_align_block_size_triton(
111
cumsum = torch.zeros((num_experts + 1, ),
112
dtype=torch.int32,
113
device=topk_ids.device)
118
- tokens_per_thread = ceil_div(numel, num_experts)
114
+ tokens_per_thread = cdiv(numel, num_experts)
115
116
moe_align_block_size_stage1[grid](
117
topk_ids,
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+3
-6
@@ -19,7 +19,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
19
CUTLASS_BLOCK_FP8_SUPPORTED)
20
from vllm.platforms import current_platform
21
from vllm.triton_utils import tl, triton
22
-from vllm.utils import direct_register_custom_op
22
+from vllm.utils import cdiv, direct_register_custom_op
23
24
logger = init_logger(__name__)
25
has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None
@@ -158,12 +158,9 @@ def apply_w8a8_block_fp8_linear(
158
if current_platform.is_cuda():
159
if current_platform.has_device_capability(100):
160
161
- def ceil_div(x: int, y: int) -> int:
162
- return (x + y - 1) // y
163
-
161
use_cutlass = cutlass_block_fp8_supported and (
165
- ceil_div(weight.shape[0], 128) == weight_scale.shape[0]
166
- and ceil_div(weight.shape[1], 128) == weight_scale.shape[1])
162
+ cdiv(weight.shape[0], 128) == weight_scale.shape[0]
163
+ and cdiv(weight.shape[1], 128) == weight_scale.shape[1])
164
else:
165
# TODO: update this after switching to public sm90 block scale gemm
166
# as it also supports weight.shape % 128 != 0