triton code pre commit for rejection_sampler.py

Signed-off-by: yuxingcyx <[email protected]>

Author: yuxingcyx

vllm_ascend/sample/rejection_sampler.py

@@ -156,7 +156,7 @@ def rejection_sample(
             BLOCK_SIZE = 2
             grid = triton.cdiv(n, BLOCK_SIZE)
             if n >= 40:
-                grid = 40 # Empirically tuned value
+                grid = 40  # Empirically tuned value
                 BLOCK_SIZE = triton.next_power_of_2(n // grid)
 
             if min(num_draft_tokens) == 1 and max(
@@ -167,7 +167,7 @@ def rejection_sample(
                     target_argmax,
                     bonus_token_ids,
                     vec_len,
-                    BLOCK_SIZE = BLOCK_SIZE,
+                    BLOCK_SIZE=BLOCK_SIZE,
                 )
             else:
                 rejection_greedy_sample_triton[(grid, )](
@@ -307,7 +307,7 @@ def expand_batch_to_tokens(
             replace_to,
             vec_len,
             MAX_NUM_TOKENS=MAX_SPEC_LEN,  # To avoid recompilation.
-            BLOCK_SIZE = BLOCK_SIZE,
+            BLOCK_SIZE=BLOCK_SIZE,
         )
     else:
         expand_pytorch(