2f03035a61 · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/_custom_ops.py

+5 -8

@@ -12,6 +12,7 @@ from vllm.scalar_type import ScalarType

12 from vllm.utils.flashinfer import (

13 flashinfer_quant_nvfp4_8x4_sf_layout,

14 )

15 +from vllm.utils.math_utils import cdiv

16

17 logger = init_logger(__name__)

18

@@ -3111,10 +3112,6 @@ def matmul_ada_mxf4_bf16_tn(

3112 return torch.ops._qutlass_C.matmul_ada_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha)

3113

3114

3114 -def ceil_div(a, b):

3115 - return (a + b - 1) // b

3116 -

3117 -

3115 if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxQuest"):

3116

3117 @register_fake("_qutlass_C::fusedQuantizeMxQuest")

@@ -3148,8 +3145,8 @@ def fusedQuantizeMx(

3145 )

3146

3147 rows, cols = a.numel() // a.size(-1), a.size(-1) // 32

3151 - n_row_blocks = ceil_div(rows, 128)

3152 - n_col_blocks = ceil_div(cols, 4)

3148 + n_row_blocks = cdiv(rows, 128)

3149 + n_col_blocks = cdiv(cols, 4)

3150 padded_rows = n_row_blocks * 128

3151 padded_cols = n_col_blocks * 4

3152

@@ -3192,8 +3189,8 @@ def fusedQuantizeNv(

3189 )

3190

3191 rows, cols = a.numel() // a.size(-1), a.size(-1) // 16

3195 - n_row_blocks = ceil_div(rows, 128)

3196 - n_col_blocks = ceil_div(cols, 4)

3192 + n_row_blocks = cdiv(rows, 128)

3193 + n_col_blocks = cdiv(cols, 4)

3194 padded_rows = n_row_blocks * 128

3195 padded_cols = n_col_blocks * 4

3196 xh_e4m3 = torch.empty(

vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py

+1 -1

@@ -32,7 +32,7 @@ def pplx_hidden_dim_scale_bytes(

32 align = 16

33

34 # For blocked per token: set to

35 - # ceil_div(hidden_dim, block_size) * sizeof(float32)

35 + # cdiv(hidden_dim, block_size) * sizeof(float32)

36 # For per-token: set to 4 * sizeof(float32) (x4 for alignment)

37 if quant_dtype is not None:

38 assert isinstance(quant_dtype, torch.dtype)

vllm/model_executor/layers/quantization/fp_quant.py

-4

@@ -248,10 +248,6 @@ class FPQuantLinearMethod(LinearMethodBase):

248 )

249

250

251 -def ceil_div(a, b):

252 - return (a + b - 1) // b

253 -

254 -

251 def fused_quantize_mx(

252 x_flat: torch.Tensor, hadamard_matrix: torch.Tensor, forward_method: str

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

vllm/model_executor/layers/quantization/qutlass_utils.py

+4 -7

@@ -17,6 +17,7 @@ import torch

17 from torch.library import wrap_triton

18

19 from vllm.triton_utils import tl, triton

20 +from vllm.utils.math_utils import cdiv

21

22

23 @triton.jit

@@ -141,10 +142,6 @@ def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor:

142 return out

143

144

144 -def ceil_div(a, b):

145 - return (a + b - 1) // b

146 -

147 -

145 def to_blocked(

146 input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = "triton"

147 ) -> torch.Tensor:

@@ -160,7 +157,7 @@ def to_blocked(

157 backend: "torch" (PyTorch path) or "triton" (Triton kernel)

158

159 Returns:

163 - Rearranged tensor of shape (32*ceil_div(H,128), 16*ceil_div(W,4))

160 + Rearranged tensor of shape (32*cdiv(H,128), 16*cdiv(W,4))

161 """

162 if backend == "triton":

163 return triton_mx_block_rearrange(input_matrix).flatten()

@@ -168,8 +165,8 @@ def to_blocked(

165 raise ValueError(f'backend must be "torch" or "triton", got {backend!r}')

166

167 rows, cols = input_matrix.shape

171 - n_row_blocks = ceil_div(rows, 128)

172 - n_col_blocks = ceil_div(cols, 4)

168 + n_row_blocks = cdiv(rows, 128)

169 + n_col_blocks = cdiv(cols, 4)

170

171 # Calculate the padded shape

172 padded_rows = n_row_blocks * 128

vllm/model_executor/model_loader/bitsandbytes_loader.py

+1 -5

@@ -22,6 +22,7 @@ from vllm.distributed import (

22 get_tensor_model_parallel_world_size,

23 )

24 from vllm.logger import init_logger

25 +from vllm.lora.utils import is_moe_model

26 from vllm.model_executor.layers.fused_moe import FusedMoE

27 from vllm.model_executor.layers.linear import (

28 LinearBase,

@@ -52,11 +53,6 @@ from vllm.utils.torch_utils import set_default_torch_dtype

53 logger = init_logger(__name__)

54

55

55 -def is_moe_model(model: torch.nn.Module) -> bool:

56 - """Checks if the model contains FusedMoE layers."""

57 - return bool(any(isinstance(module, FusedMoE) for module in model.modules()))

58 -

59 -

56 class BitsAndBytesModelLoader(BaseModelLoader):

57 """Model loader to load model weights with BitsAndBytes quantization."""

58

vllm/model_executor/models/aya_vision.py

+3 -8

@@ -40,6 +40,7 @@ from .siglip import SiglipVisionModel

40 from .utils import (

41 AutoWeightsLoader,

42 WeightsMapper,

43 + get_layer_index,

44 init_vllm_registered_model,

45 maybe_prefix,

46 )

@@ -300,21 +301,15 @@ def _get_num_hidden_layers(hf_config: AyaVisionConfig) -> int:

301 num_hidden_layers = hf_config.vision_config.num_hidden_layers

302 # If we have one feature layer, initialize up to that layer

303 if isinstance(feature_layers, int):

303 - return _get_layer_index(feature_layers, num_hidden_layers)

304 + return get_layer_index(feature_layers, num_hidden_layers)

305 # If we have multiple feature layers, initialize up to the deepest m

306 elif isinstance(feature_layers, (list, tuple)):

306 - return max(_get_layer_index(idx, num_hidden_layers) for idx in feature_layers)

307 + return max(get_layer_index(idx, num_hidden_layers) for idx in feature_layers)

308 raise TypeError(

309 f"vision_layer_feature type: {type(feature_layers)} is not supported"

310 )

311

312

312 -def _get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:

313 - if feature_layer_index < 0:

314 - return num_hidden_layers + feature_layer_index + 1

315 - return feature_layer_index

316 -

317 -

313 @MULTIMODAL_REGISTRY.register_processor(

314 AyaVisionMultiModalProcessor,

315 info=AyaVisionProcessingInfo,

vllm/model_executor/models/llava.py

+3 -16

@@ -63,6 +63,7 @@ from .siglip import SiglipVisionModel

63 from .utils import (

64 AutoWeightsLoader,

65 WeightsMapper,

66 + get_layer_index,

67 init_vllm_registered_model,

68 maybe_prefix,

69 )

@@ -442,29 +443,15 @@ def _get_num_hidden_layers(hf_config: LlavaLikeConfig) -> int:

443 num_hidden_layers = hf_config.vision_config.num_hidden_layers

444 # If we have one feature layer, initialize up to that layer

445 if isinstance(feature_layers, int):

445 - return _get_layer_index(feature_layers, num_hidden_layers)

446 + return get_layer_index(feature_layers, num_hidden_layers)

447 # If we have multiple feature layers, initialize up to the deepest one

448 elif isinstance(feature_layers, (list, tuple)):

448 - return max(_get_layer_index(idx, num_hidden_layers) for idx in feature_layers)

449 + return max(get_layer_index(idx, num_hidden_layers) for idx in feature_layers)

450 raise TypeError(

451 f"vision_layer_feature type: {type(feature_layers)} is not supported"

452 )

453

454

454 -def _get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:

455 - """Given a signed vision feature layer, get the number of hidden layers

456 - needed to leverage it.

457 -

458 - Args:

459 - feature_layer_index: Index of a required layer in the visual encoder.

460 - num_hidden_layers: The total number of hidden layers in the visual

461 - encoder.

462 - """

463 - if feature_layer_index < 0:

464 - return num_hidden_layers + feature_layer_index + 1

465 - return feature_layer_index

466 -

467 -

455 def init_vision_tower_for_llava(

456 hf_config: LlavaLikeConfig,

457 quant_config: QuantizationConfig | None,

vllm/model_executor/models/mistral3.py

+3 -16

@@ -52,6 +52,7 @@ from .pixtral import PixtralHFEncoderInfo, PixtralHFVisionModel

52 from .utils import (

53 AutoWeightsLoader,

54 WeightsMapper,

55 + get_layer_index,

56 init_vllm_registered_model,

57 maybe_prefix,

58 )

@@ -369,29 +370,15 @@ def _get_num_hidden_layers(hf_config: LlavaLikeConfig) -> int:

370 num_hidden_layers = hf_config.vision_config.num_hidden_layers

371 # If we have one feature layer, initialize up to that layer

372 if isinstance(feature_layers, int):

372 - return _get_layer_index(feature_layers, num_hidden_layers)

373 + return get_layer_index(feature_layers, num_hidden_layers)

374 # If we have multiple feature layers, initialize up to the deepest one

375 elif isinstance(feature_layers, (list, tuple)):

375 - return max(_get_layer_index(idx, num_hidden_layers) for idx in feature_layers)

376 + return max(get_layer_index(idx, num_hidden_layers) for idx in feature_layers)

377 raise TypeError(

378 f"vision_layer_feature type: {type(feature_layers)} is not supported"

379 )

380

381

381 -def _get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:

382 - """Given a signed vision feature layer, get the number of hidden layers

383 - needed to leverage it.

384 -

385 - Args:

386 - feature_layer_index: Index of a required layer in the visual encoder.

387 - num_hidden_layers: The total number of hidden layers in the visual

388 - encoder.

389 - """

390 - if feature_layer_index < 0:

391 - return num_hidden_layers + feature_layer_index + 1

392 - return feature_layer_index

393 -

394 -

382 def init_vision_tower_for_llava(

383 hf_config: LlavaLikeConfig,

384 quant_config: QuantizationConfig | None,

vllm/model_executor/models/tarsier.py

+8 -8

@@ -47,7 +47,12 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

47 from .clip import CLIPVisionModel

48 from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

49 from .siglip import SiglipVisionModel

50 -from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix

50 +from .utils import (

51 + AutoWeightsLoader,

52 + get_layer_index,

53 + init_vllm_registered_model,

54 + maybe_prefix,

55 +)

56 from .vision import (

57 VisionEncoderInfo,

58 get_num_selected_vision_tokens,

@@ -356,18 +361,13 @@ def init_vision_tower_for_tarsier(

361 feature_layers = hf_config.vision_feature_layer

362 base_num_hidden_layers = vision_config.num_hidden_layers

363

359 - def _get_layer_index(feature_layer_index: int, num_hidden_layers_total: int) -> int:

360 - if feature_layer_index < 0:

361 - return num_hidden_layers_total + feature_layer_index + 1

362 - return feature_layer_index

363 -

364 if isinstance(feature_layers, int):

365 - num_hidden_layers_to_init = _get_layer_index(

365 + num_hidden_layers_to_init = get_layer_index(

366 feature_layers, base_num_hidden_layers

367 )

368 elif isinstance(feature_layers, (list, tuple)):

369 num_hidden_layers_to_init = max(

370 - _get_layer_index(idx, base_num_hidden_layers) for idx in feature_layers

370 + get_layer_index(idx, base_num_hidden_layers) for idx in feature_layers

371 )

372 else:

373 raise TypeError(

vllm/model_executor/models/utils.py

+13

@@ -830,3 +830,16 @@ def process_eagle_weight(

830 model.has_own_lm_head = True

831 if "embed_tokens" in name:

832 model.has_own_embed_tokens = True

833 +

834 +

835 +def get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:

836 + """Given a signed vision feature layer, get the number of hidden layers

837 + needed to leverage it.

838 +

839 + Args:

840 + feature_layer_index: Index of a required layer in the visual encoder.

841 + num_hidden_layers: The total number of hidden layers in the visual encoder.

842 + """

843 + if feature_layer_index < 0:

844 + return num_hidden_layers + feature_layer_index + 1

845 + return feature_layer_index

vllm/model_executor/warmup/deep_gemm_warmup.py

+2 -4

@@ -25,6 +25,7 @@ from vllm.utils.deep_gemm import (

25 get_mk_alignment_for_contiguous_layout,

26 m_grouped_fp8_gemm_nt_contiguous,

27 )

28 +from vllm.utils.math_utils import cdiv

29

30

31 def _generate_optimal_warmup_m_values(

@@ -40,9 +41,6 @@ def _generate_optimal_warmup_m_values(

41 device: The torch device to get properties from.

42 """

43

43 - def ceil_div(a: int, b: int) -> int:

44 - return (a + b - 1) // b

45 -

44 # DeepGEMM's possible block sizes

45 block_ms = [64, 128, 256]

46 block_ns = list(range(16, min(257, n + 1), 16))

@@ -63,7 +61,7 @@ def _generate_optimal_warmup_m_values(

61 for wave in range(1, 11): # Up to 10 waves

62 # M where this block config transitions to next wave

63 target_blocks = wave * num_sms

66 - m = target_blocks * block_m // ceil_div(n, block_n)

64 + m = target_blocks * block_m // cdiv(n, block_n)

65 if 1 <= m <= max_tokens:

66 m_values.add(m)

67

"refactor: refactor_repeated_interfaces" (#32486)