Update deprecated Python 3.8 typing (#13971)

benchmarks/kernels/benchmark_machete.py

@@ -7,9 +7,10 @@ import math
 import os
 import pickle as pkl
 import time
+from collections.abc import Iterable
 from dataclasses import dataclass
 from itertools import product
-from typing import Callable, Iterable, List, Optional, Tuple
+from typing import Callable, Optional
 
 import pandas as pd
 import torch
@@ -102,8 +103,8 @@ def quantize_and_pack(atype: torch.dtype,
     return w_ref, w_q, w_s, w_zp
 
 
-def create_bench_tensors(shape: Tuple[int, int, int], types: TypeConfig,
-                         group_size: Optional[int]) -> List[BenchmarkTensors]:
+def create_bench_tensors(shape: tuple[int, int, int], types: TypeConfig,
+                         group_size: Optional[int]) -> list[BenchmarkTensors]:
     m, n, k = shape
 
     # we want to make sure that weights don't fit into L2 cache between runs so
@@ -114,7 +115,7 @@ def create_bench_tensors(shape: Tuple[int, int, int], types: TypeConfig,
 
     a = rand_data((m, k), types.act_type, scale=5)
 
-    benchmark_tensors: List[BenchmarkTensors] = []
+    benchmark_tensors: list[BenchmarkTensors] = []
     for _ in range(num_weights):
         w = rand_data((k, n), types.act_type, scale=5)
 
@@ -276,7 +277,7 @@ def machete_create_bench_fn(bt: BenchmarkTensors,
 
 
 def bench_fns(label: str, sub_label: str, description: str,
-              fns: List[Callable]):
+              fns: list[Callable]):
 
     min_run_time = 1 if not NVTX_PROFILE else 0.1
     res = TBenchmark.Timer(
@@ -311,7 +312,7 @@ def bench(types: TypeConfig,
           n: int,
           label: str,
           sub_label: str,
-          sweep_schedules: bool = True) -> List[TMeasurement]:
+          sweep_schedules: bool = True) -> list[TMeasurement]:
     benchmark_tensors = create_bench_tensors((m, n, k), types, group_size)
     sub_label += f", L={len(benchmark_tensors)}"
 
@@ -414,12 +415,12 @@ def bench(types: TypeConfig,
 
 
 # runner
-def print_timers(timers: List[TMeasurement]):
+def print_timers(timers: list[TMeasurement]):
     compare = TBenchmark.Compare(timers)
     compare.print()
 
 
-def run(args, MKNs: Iterable[Tuple[int, int, int]]) -> Iterable[TMeasurement]:
+def run(args, MKNs: Iterable[tuple[int, int, int]]) -> Iterable[TMeasurement]:
     types = TypeConfig(
         act_type=args.act_type,
         weight_type=scalar_types.uint4b8 if args.group_zero_type is None \
@@ -431,7 +432,7 @@ def run(args, MKNs: Iterable[Tuple[int, int, int]]) -> Iterable[TMeasurement]:
         token_scale_type=args.token_scale_type,
     )
 
-    results: List[TMeasurement] = []
+    results: list[TMeasurement] = []
     for m, k, n in MKNs:
         timers = bench(types,
                        args.group_size,
@@ -449,8 +450,8 @@ def run(args, MKNs: Iterable[Tuple[int, int, int]]) -> Iterable[TMeasurement]:
 
 # output makers
 def make_output(
-    data: List[TMeasurement],
-    MKNs: Iterable[Tuple[int, int, int]],
+    data: list[TMeasurement],
+    MKNs: Iterable[tuple[int, int, int]],
     base_description: str,
     timestamp=None,
 ):
@@ -497,7 +498,7 @@ def run_model_bench(args):
     for i, model in enumerate(args.models):
         print(f"[{i}]  {model}")
 
-    def model_shapes(model_name: str, tp_size: int) -> List[Tuple[int, int]]:
+    def model_shapes(model_name: str, tp_size: int) -> list[tuple[int, int]]:
         KNs = []
         for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]):
             KN[tp_split_dim] = KN[tp_split_dim] // tp_size