b893d661b1 · tduncan/Nvidia-Optimization-Stack-VLLM

pyproject.toml

-34

@@ -101,40 +101,6 @@ include = ["vllm*"]

101 "vllm/v1/engine/utils.py" = ["E501"]

102 "vllm/v1/utils.py" = ["E501"]

103 "vllm/v1/worker/gpu_model_runner.py" = ["E501"]

104 -## Simplification rules

105 -"tests/distributed/test_expert_placement.py" = ["SIM108"]

106 -"tests/kernels/attention/test_cutlass_mla_decode.py" = ["SIM108"]

107 -"tests/kernels/attention/test_flashmla.py" = ["SIM108"]

108 -"tests/kernels/attention/test_lightning_attn.py" = ["SIM108"]

109 -"tests/kernels/moe/test_pplx_moe.py" = ["SIM108"]

110 -"tests/kernels/quantization/test_cutlass_scaled_mm.py" = ["SIM108"]

111 -"tests/kernels/test_onednn.py" = ["SIM108"]

112 -"tests/kernels/utils.py" = ["SIM108"]

113 -"tests/multimodal/test_processing.py" = ["SIM108"]

114 -"vllm/attention/ops/triton_reshape_and_cache_flash.py" = ["SIM108"]

115 -"vllm/distributed/parallel_state.py" = ["SIM108"]

116 -"vllm/entrypoints/chat_utils.py" = ["SIM108"]

117 -"vllm/entrypoints/llm.py" = ["SIM108"]

118 -"vllm/executor/ray_distributed_executor.py" = ["SIM108", "SIM112"]

119 -"vllm/model_executor/layers/batch_invariant.py" = ["SIM108"]

120 -"vllm/model_executor/layers/fla/ops/chunk_o.py" = ["SIM108"]

121 -"vllm/model_executor/layers/fused_moe/fused_moe.py" = ["SIM108"]

122 -"vllm/model_executor/layers/fused_moe/layer.py" = ["SIM108"]

123 -"vllm/model_executor/layers/fused_moe/modular_kernel.py" = ["SIM108"]

124 -"vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py" = ["SIM108"]

125 -"vllm/model_executor/layers/layernorm.py" = ["SIM108"]

126 -"vllm/model_executor/layers/lightning_attn.py" = ["SIM108"]

127 -"vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py" = ["SIM103"]

128 -"vllm/model_executor/layers/quantization/compressed_tensors/utils.py" = ["SIM110"]

129 -"vllm/model_executor/layers/quantization/quark/utils.py" = ["SIM110"]

130 -"vllm/utils/__init__.py" = ["SIM108"]

131 -"vllm/v1/sample/ops/bad_words.py" = ["SIM108"]

132 -"vllm/v1/sample/rejection_sampler.py" = ["SIM108"]

133 -"vllm/v1/worker/tpu_model_runner.py" = ["SIM108"]

134 -"vllm/_custom_ops.py" = ["SIM108"]

135 -"tools/profiler/print_layerwise_table.py" = ["SIM118"]

136 -## Loop variable binding issues

137 -"tests/kernels/mamba/test_mamba_ssm_ssd.py" = ["B023"]

104 # End of temporary ignores

105

106 [tool.ruff.lint]

tests/distributed/test_expert_placement.py

+1 -4

@@ -12,10 +12,7 @@ def verify_round_robin_pattern(expert_map, ep_rank, ep_size, global_num_experts)

12 base_experts = global_num_experts // ep_size

13 remainder = global_num_experts % ep_size

14

15 - if ep_rank < remainder:

16 - local_num_experts = base_experts + 1

17 - else:

18 - local_num_experts = base_experts

15 + local_num_experts = base_experts + 1 if ep_rank < remainder else base_experts

16

17 # Expected expert IDs for this rank in round_robin pattern

18 # For non-divisible cases, ranks with extra experts start earlier

tests/kernels/attention/test_cutlass_mla_decode.py

+1 -4

@@ -66,10 +66,7 @@ def test_cutlass_mla_decode(

66 b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal, varlen, torch_dtype

67 ):

68 device = torch.device("cuda:0")

69 - if torch_dtype == torch.float8_e4m3fn:

70 - init_dtype = torch.bfloat16

71 - else:

72 - init_dtype = torch_dtype

69 + init_dtype = torch.bfloat16 if torch_dtype == torch.float8_e4m3fn else torch_dtype

70 torch.set_default_dtype(init_dtype)

71 torch.set_default_device(device)

72 torch.cuda.set_device(device)

tests/kernels/attention/test_flashmla.py

+1 -4

@@ -52,10 +52,7 @@ def test_flash_mla(

52 b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal, varlen, torch_dtype

53 ):

54 device = torch.device("cuda:0")

55 - if torch_dtype == torch.float8_e4m3fn:

56 - init_dtype = torch.bfloat16

57 - else:

58 - init_dtype = torch_dtype

55 + init_dtype = torch.bfloat16 if torch_dtype == torch.float8_e4m3fn else torch_dtype

56 torch.set_default_dtype(init_dtype)

57 torch.set_default_device(device)

58 torch.cuda.set_device(device)

tests/kernels/attention/test_lightning_attn.py

+1 -4

@@ -33,10 +33,7 @@ def reference_lightning_attention(q, k, v, ed, block_size, kv_history):

33

34 # More efficient implementation

35 # Convert decay factors to matrix form

36 - if ed.dim() == 1:

37 - decay = torch.exp(-ed).view(1, -1, 1, 1)

38 - else:

39 - decay = torch.exp(-ed)

36 + decay = torch.exp(-ed).view(1, -1, 1, 1) if ed.dim() == 1 else torch.exp(-ed)

37

38 for b in range(B):

39 for step in range(S):

tests/kernels/moe/test_pplx_moe.py

+1 -4

@@ -705,10 +705,7 @@ def _pplx_moe(

705 with set_current_vllm_config(vllm_config), override_config(moe_config):

706 topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)

707

708 - if shared_experts is not None:

709 - shared_output = shared_experts(a)

710 - else:

711 - shared_output = None

708 + shared_output = shared_experts(a) if shared_experts is not None else None

709

710 torch_output = torch_experts(

711 a,

tests/kernels/quantization/test_cutlass_scaled_mm.py

+2 -8

@@ -88,10 +88,7 @@ def cutlass_fp8_gemm_helper(

88 # make scales K-major for blockwise quant, doesn't affect 1D scales

89 scale_b = scale_b.t().contiguous().t()

90

91 - if use_bias:

92 - bias = torch.rand((n,), device=device, dtype=out_dtype) * 10

93 - else:

94 - bias = None

91 + bias = torch.rand((n,), device=device, dtype=out_dtype) * 10 if use_bias else None

92

93 out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)

94 baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)

@@ -122,10 +119,7 @@ def cutlass_int8_gemm_helper(

119 scale_a = torch.randn(a_scales_shape, device=device, dtype=torch.float32)

120 scale_b = torch.randn(b_scales_shape, device=device, dtype=torch.float32)

121

125 - if use_bias:

126 - bias = torch.rand((n,), device=device, dtype=out_dtype) * 10

127 - else:

128 - bias = None

122 + bias = torch.rand((n,), device=device, dtype=out_dtype) * 10 if use_bias else None

123

124 out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)

125 baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)

tests/kernels/test_onednn.py

+1 -4

@@ -84,10 +84,7 @@ def onednn_int8_gemm_test_helper(

84 azp = None

85 azp_adj = None

86

87 - if use_bias:

88 - bias = torch.rand((n,), device=device, dtype=out_dtype) * 10

89 - else:

90 - bias = None

87 + bias = torch.rand((n,), device=device, dtype=out_dtype) * 10 if use_bias else None

88

89 handler = ops.create_onednn_scaled_mm(

90 b,

tests/kernels/utils.py

+3 -7

@@ -963,13 +963,9 @@ def make_test_metadata(

963 None if encoder_seq_lens is None else (sum(encoder_seq_lens))

964 )

965

966 - if cross_test_params is None:

967 - cross_kv_mmap = None

968 - else:

969 - # Encoder/decoder or encoder-only models only:

970 - # * Extract *cross-attention* slot_mapping and block table

971 - # (kv_mmap)

972 - cross_kv_mmap = cross_test_params.kv_mmap

966 + # For encoder/decoder or encoder-only models only, extract *cross-attention*

967 + # slot_mapping and block table (kv_mmap)

968 + cross_kv_mmap = None if cross_test_params is None else cross_test_params.kv_mmap

969

970 attn_backend_obj = make_backend(attn_backend.name)

971

tests/multimodal/test_processing.py

+2 -8

@@ -941,10 +941,7 @@ def test_limit_mm_per_prompt_dummy(model_id, limit, num_supported, is_valid):

941

942 profiler = MultiModalProfiler(processor)

943

944 - if is_valid:

945 - exc_ctx = nullcontext()

946 - else:

947 - exc_ctx = pytest.raises(ValueError, match="At most")

944 + exc_ctx = nullcontext() if is_valid else pytest.raises(ValueError, match="At most")

945

946 with exc_ctx:

947 profiler.get_decoder_dummy_data(

@@ -985,10 +982,7 @@ def test_limit_mm_per_prompt_apply(model_id, num_images, limit, is_valid):

982 else:

983 mm_data = {"image": [image] * num_images}

984

988 - if is_valid:

989 - exc_ctx = nullcontext()

990 - else:

991 - exc_ctx = pytest.raises(ValueError, match="At most")

985 + exc_ctx = nullcontext() if is_valid else pytest.raises(ValueError, match="At most")

986

987 with exc_ctx:

988 processor.apply(

tools/profiler/print_layerwise_table.py

+1 -1

@@ -58,7 +58,7 @@ if __name__ == "__main__":

58

59 assert args.phase in profile_data, (

60 f"Cannot find phase {args.phase} in profile data. Choose one among"

61 - f"{[x for x in profile_data.keys() if 'prefill' in x or 'decode' in x]}"

61 + f"{[x for x in profile_data if 'prefill' in x or 'decode' in x]}"

62 ) # noqa

63

64 if args.table == "summary":

vllm/_custom_ops.py

+1 -4

@@ -2370,10 +2370,7 @@ class CPUDNNLGEMMHandler:

2370 torch.ops._C.release_dnnl_matmul_handler(self.handler)

2371

2372

2373 -if hasattr(torch.ops._C, "create_onednn_mm_handler"):

2374 - _supports_onednn = True

2375 -else:

2376 - _supports_onednn = False

2373 +_supports_onednn = bool(hasattr(torch.ops._C, "create_onednn_mm_handler"))

2374

2375

2376 def is_onednn_acl_supported():

vllm/attention/ops/triton_reshape_and_cache_flash.py

+3 -6

@@ -52,12 +52,9 @@ def reshape_and_cache_kernel_flash(

52 key_ptr + src_key_idx + tile_pos, mask=tile_pos < (num_heads * head_size)

53 )

54 if FP8_KV_CACHE:

55 - if key_load.dtype.is_fp8():

56 - key_tile = key_load

57 - else:

58 - # tl.store will do the correct implicit cast to fp8,

59 - # based on the key_cache_ptr.dtype.element_ty

60 - key_tile = key_load / tl.load(k_scale)

55 + # tl.store will do the correct implicit cast to fp8,

56 + # based on the key_cache_ptr.dtype.element_ty

57 + key_tile = key_load if key_load.dtype.is_fp8() else key_load / tl.load(k_scale)

58 else:

59 key_tile = key_load

60

vllm/distributed/parallel_state.py

+1 -4

@@ -1097,10 +1097,7 @@ def init_distributed_environment(

1097 if local_rank == -1:

1098 # local rank not set, this usually happens in single-node

1099 # setting, where we can use rank as local rank

1100 - if distributed_init_method == "env://":

1101 - local_rank = envs.LOCAL_RANK

1102 - else:

1103 - local_rank = rank

1100 + local_rank = envs.LOCAL_RANK if distributed_init_method == "env://" else rank

1101 global _WORLD, _NODE_COUNT

1102 if _WORLD is None:

1103 ranks = list(range(torch.distributed.get_world_size()))

vllm/entrypoints/chat_utils.py

+1 -4

@@ -1310,10 +1310,7 @@ def _parse_chat_message_content_part(

1310

1311 modality = None

1312 if part_type == "image_pil":

1313 - if content is not None:

1314 - image_content = cast(Image.Image, content)

1315 - else:

1316 - image_content = None

1313 + image_content = cast(Image.Image, content) if content is not None else None

1314 mm_parser.parse_image_pil(image_content, uuid)

1315 modality = "image"

1316 elif part_type in ("image_url", "input_image"):

vllm/entrypoints/llm.py

+1 -4

@@ -1018,10 +1018,7 @@ class LLM:

1018 pooling_task = "encode"

1019

1020 if pooling_task is None:

1021 - if "embed" in self.supported_tasks:

1022 - pooling_task = "embed"

1023 - else:

1024 - pooling_task = "encode"

1021 + pooling_task = "embed" if "embed" in self.supported_tasks else "encode"

1022

1023 logger.warning_once(

1024 "`LLM.encode` is currently using `pooling_task = %s`.\n"

vllm/executor/ray_distributed_executor.py

+5 -10

@@ -458,10 +458,7 @@ class RayDistributedExecutor(DistributedExecutorBase):

458 else:

459 serialized_data = self.input_encoder.encode(execute_model_req)

460 outputs = ray.get(self.forward_dag.execute(serialized_data))

461 - if self.use_v1:

462 - output = outputs[0]

463 - else:

464 - output = self.output_decoder.decode(outputs[0])

461 + output = outputs[0] if self.use_v1 else self.output_decoder.decode(outputs[0])

462 return output

463

464 def _run_workers(

@@ -482,10 +479,7 @@ class RayDistributedExecutor(DistributedExecutorBase):

479 rather than blocking on the results.

480 - args/kwargs: All workers share the same args/kwargs

481 """

485 - if isinstance(method, str):

486 - sent_method = method

487 - else:

488 - sent_method = cloudpickle.dumps(method)

482 + sent_method = method if isinstance(method, str) else cloudpickle.dumps(method)

483 del method

484 if self.use_ray_spmd_worker:

485 assert not async_run_tensor_parallel_workers_only, (

@@ -573,8 +567,9 @@ class RayDistributedExecutor(DistributedExecutorBase):

567 from ray.dag import InputNode, MultiOutputNode

568

569 logger.info(

576 - "RAY_CGRAPH_get_timeout is set to %s", os.environ["RAY_CGRAPH_get_timeout"]

577 - ) # noqa: SIM112

570 + "RAY_CGRAPH_get_timeout is set to %s",

571 + os.environ["RAY_CGRAPH_get_timeout"], # noqa: SIM112

572 + )

573 logger.info(

574 "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE = %s",

575 envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE,

vllm/model_executor/layers/batch_invariant.py

+1 -4

@@ -439,10 +439,7 @@ def mean_dim(

439 output = torch.empty(output_shape, dtype=dtype, device=input.device)

440

441 # Reshape output for kernel

442 - if keepdim:

443 - output_2d = output.reshape(M, 1, K).squeeze(1)

444 - else:

445 - output_2d = output.reshape(M, K)

442 + output_2d = output.reshape(M, 1, K).squeeze(1) if keepdim else output.reshape(M, K)

443

444 # Launch kernel

445 grid = (M * K,)

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

+1 -4

@@ -151,10 +151,7 @@ def chunk_fwd_o(

151 ) -> torch.Tensor:

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

153 H = v.shape[-2]

154 - if FLA_GDN_FIX_BT:

155 - BT = 64

156 - else:

157 - BT = min(chunk_size, max(16, triton.next_power_of_2(T)))

154 + BT = 64 if FLA_GDN_FIX_BT else min(chunk_size, max(16, triton.next_power_of_2(T)))

155 chunk_indices = (

156 prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None

157 )

vllm/model_executor/layers/fused_moe/fused_moe.py

+1 -4

@@ -1746,10 +1746,7 @@ def fused_experts_impl(

1746 else:

1747 raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")

1748

1749 - if inplace:

1750 - out_hidden_states = hidden_states

1751 - else:

1752 - out_hidden_states = torch.empty_like(hidden_states)

1749 + out_hidden_states = hidden_states if inplace else torch.empty_like(hidden_states)

1750

1751 if use_mxfp4_w4a4:

1752 # Weight has to be dequantized for mxfp4 emulation.

vllm/model_executor/layers/fused_moe/layer.py

+1 -4

@@ -886,10 +886,7 @@ def determine_expert_map(

886 # Distribute experts as evenly as possible to each rank.

887 base_experts = global_num_experts // ep_size

888 remainder = global_num_experts % ep_size

889 - if ep_rank < remainder:

890 - local_num_experts = base_experts + 1

891 - else:

892 - local_num_experts = base_experts

889 + local_num_experts = base_experts + 1 if ep_rank < remainder else base_experts

890

891 # Create a tensor of size num_experts filled with -1

892 expert_map = torch.full((global_num_experts,), -1, dtype=torch.int32)

vllm/model_executor/layers/fused_moe/modular_kernel.py

+1 -4

@@ -948,10 +948,7 @@ class FusedMoEModularKernel(torch.nn.Module):

948 """

949

950 a1 = hidden_states

951 - if inplace and self.shared_experts is None:

952 - output = a1

953 - else:

954 - output = torch.zeros_like(a1)

951 + output = a1 if inplace and self.shared_experts is None else torch.zeros_like(a1)

952

953 local_num_experts = w1.size(0)

954 if global_num_experts == -1:

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

+1 -4

@@ -355,10 +355,7 @@ def rocm_aiter_fused_experts(

355 topk_weights = topk_weights.to(torch.float32)

356 topk_ids = topk_ids.to(torch.int32)

357

358 - if expert_map is not None:

359 - expert_mask = (expert_map > -1).to(torch.int32)

360 - else:

361 - expert_mask = None

358 + expert_mask = (expert_map > -1).to(torch.int32) if expert_map is not None else None

359

360 # w8a8 per-channel quantization

361 if (

vllm/model_executor/layers/layernorm.py

+1 -4

@@ -318,10 +318,7 @@ class GemmaRMSNorm(CustomOp):

318 """PyTorch-native implementation equivalent to forward()."""

319 orig_dtype = x.dtype

320 if residual is not None:

321 - if orig_dtype == torch.float16:

322 - x = x + residual.float()

323 - else:

324 - x = x + residual

321 + x = x + residual.float() if orig_dtype == torch.float16 else x + residual

322 residual = x

323

324 x = x.float()

vllm/model_executor/layers/lightning_attn.py

+1 -4

@@ -207,10 +207,7 @@ def _fwd_kv_parallel(

207 kv = tl.zeros([D_FBLOCK, E_FBLOCK], dtype=tl.float32)

208

209 # Handle the last block which might be smaller than BLOCK

210 - if off_block == NUM_BLOCK - 1:

211 - split_n = n - (NUM_BLOCK - 1) * BLOCK

212 - else:

213 - split_n = BLOCK

210 + split_n = n - (NUM_BLOCK - 1) * BLOCK if off_block == NUM_BLOCK - 1 else BLOCK

211 left_shift = tl.cdiv(split_n, CBLOCK) * CBLOCK - split_n

212 num_blocks = min(tl.cdiv(split_n, CBLOCK), NUM_CBLOCK)

213 k_decay_ptr += (NUM_CBLOCK - num_blocks) * CBLOCK

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py

+3 -7

@@ -502,15 +502,11 @@ class CompressedTensorsConfig(QuantizationConfig):

502 QuantizationStrategy.CHANNEL,

503 QuantizationStrategy.BLOCK,

504 ]

505 - if not (

505 + return (

506 is_symmetric_weight

507 - and is_static_weight # noqa: SIM103

507 + and is_static_weight

508 and is_tensor_or_channel_or_block_weight

509 - ):

510 - return False

511 -

512 - # All conditions satisfied.

513 - return True

509 + )

510

511 def _is_wNa16_group_channel(

512 self, weight_quant: QuantizationArgs, input_quant: QuantizationArgs

vllm/model_executor/layers/quantization/compressed_tensors/utils.py

+1 -4

@@ -80,10 +80,7 @@ def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool:

80 Checks whether a layer_name is exactly equal or a regex match for

81 if target starts with 're:' to any target in list.

82 """

83 - for target in targets:

84 - if _is_equal_or_regex_match(layer_name, target):

85 - return True

86 - return False

83 + return any(_is_equal_or_regex_match(layer_name, target) for target in targets)

84

85

86 def find_matched_target(

vllm/model_executor/layers/quantization/quark/utils.py

+1 -4

@@ -81,10 +81,7 @@ def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool:

81 Checks whether a layer_name is exactly equal or a regex match for

82 if target starts with 're:' to any target in list.

83 """

84 - for target in targets:

85 - if _is_equal_or_regex_match(layer_name, target):

86 - return True

87 - return False

84 + return any(_is_equal_or_regex_match(layer_name, target) for target in targets)

85

86

87 def _is_equal_or_regex_match(

vllm/utils/__init__.py

+1 -4

@@ -3052,10 +3052,7 @@ def make_zmq_socket(

3052 # - Set a large 0.5GB buffer to improve throughput

3053 # For systems with less memory:

3054 # - Use system default (-1) to avoid excessive memory consumption

3055 - if total_mem > 32 and available_mem > 16:

3056 - buf_size = int(0.5 * 1024**3) # 0.5GB in bytes

3057 - else:

3058 - buf_size = -1 # Use system default buffer size

3055 + buf_size = int(0.5 * 1024**3) if total_mem > 32 and available_mem > 16 else -1

3056

3057 if bind is None:

3058 bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB)

vllm/v1/sample/ops/bad_words.py

+1 -4

@@ -17,10 +17,7 @@ def _apply_bad_words_single_batch(

17

18 prefix_length = len(bad_word_ids) - 1

19 last_token_id = bad_word_ids[-1]

20 - if prefix_length > 0:

21 - actual_prefix = past_tokens_ids[-prefix_length:]

22 - else:

23 - actual_prefix = []

20 + actual_prefix = past_tokens_ids[-prefix_length:] if prefix_length > 0 else []

21 expected_prefix = bad_word_ids[:prefix_length]

22

23 assert len(actual_prefix) == len(expected_prefix)

vllm/v1/sample/rejection_sampler.py

+4 -16

@@ -444,18 +444,12 @@ def rejection_greedy_sample_kernel(

444 req_idx = tl.program_id(0)

445 # FIXME(woosuk): Because is_greedy_ptr is not None at profiling run,

446 # re-compilation may happen during runtime when is_greedy_ptr is None.

447 - if is_greedy_ptr is None:

448 - is_greedy = True

449 - else:

450 - is_greedy = tl.load(is_greedy_ptr + req_idx)

447 + is_greedy = True if is_greedy_ptr is None else tl.load(is_greedy_ptr + req_idx)

448 if not is_greedy:

449 # Early exit for non-greedy sampling requests.

450 return

451

455 - if req_idx == 0:

456 - start_idx = 0

457 - else:

458 - start_idx = tl.load(cu_num_draft_tokens_ptr + req_idx - 1)

452 + start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)

453 end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)

454 num_draft_tokens = end_idx - start_idx

455

@@ -503,10 +497,7 @@ def rejection_random_sample_kernel(

497 # Early exit for greedy sampling requests.

498 return

499

506 - if req_idx == 0:

507 - start_idx = 0

508 - else:

509 - start_idx = tl.load(cu_num_draft_tokens_ptr + req_idx - 1)

500 + start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)

501 end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)

502 num_draft_tokens = end_idx - start_idx

503

@@ -583,10 +574,7 @@ def sample_recovered_tokens_kernel(

574 NO_DRAFT_PROBS: tl.constexpr,

575 ):

576 req_idx = tl.program_id(0)

586 - if req_idx == 0:

587 - start_idx = 0

588 - else:

589 - start_idx = tl.load(cu_num_draft_tokens_ptr + req_idx - 1)

577 + start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)

578 end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)

579 num_draft_tokens = end_idx - start_idx

580

vllm/v1/worker/tpu_model_runner.py

+2 -6

@@ -507,12 +507,8 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

507 removed_req_indices = sorted(removed_req_indices, reverse=True)

508 for req_id in req_ids_to_add:

509 req_state = self.requests[req_id]

510 - if removed_req_indices:

511 - # Fill the empty index.

512 - req_index = removed_req_indices.pop()

513 - else:

514 - # Append to the end.

515 - req_index = None

510 + # Fill the empty index or append to the end

511 + req_index = removed_req_indices.pop() if removed_req_indices else None

512 self.input_batch.add_request(req_state, req_index)

513

514 # Condense the batched states if there are empty indices.

Fix per file ruff ignores related to simplification (#26259)