docs(python): condense verbose comments (comments-only, no code change)

aiter/__init__.py

@@ -58,9 +58,8 @@ def getLogger():
 
 logger = getLogger()
 AITER_AOT_IMPORT = os.getenv("AITER_AOT_IMPORT", "0") == "1"
-# Triton-only: expose only the Triton ops, skipping the C++/CK/HIP ops and their
-# JIT build. Always on for Windows (no CK/HIP there); elsewhere opt in via the
-# env var, e.g. Triton-backend users with no C++ toolchain or CK.
+# Triton-only: expose only the Triton ops, skipping the C++/CK/HIP ops and their JIT build. Always on for Windows (no
+# CK/HIP there); elsewhere opt in via the env var, e.g. Triton-backend users with no C++ toolchain or CK.
 AITER_TRITON_ONLY = (
     os.getenv("AITER_TRITON_ONLY", "0") == "1" or sys.platform == "win32"
 )

aiter/aot/flydsl/chunk_gdn_h.py

@@ -66,14 +66,12 @@
 )
 DEFAULT_CSVS = [str(_DEFAULT_CSV)]
 CHUNK_GDN_H_AOT_ARCH_DEFAULT = "gfx950"
-# K5 ships a single kernel today; mirrors the ``@flyc.kernel(name=...)``
-# decorator in ``kernels/chunk_gated_delta_h.py`` so the AOT ``result`` dict
-# and any failure-mode print share one source of truth.
+# K5 ships a single kernel today; mirrors the ``@flyc.kernel(name=...)`` decorator in ``kernels/chunk_gated_delta_h.py``
+# so the AOT ``result`` dict and any failure-mode print share one source of truth.
 _KERNEL_NAME = "chunk_gdn_fwd_h_flydsl_vk"
 
 # Map jsonl ``dtype`` string -> torch dtype name used for dummy tensors.
-# Only bf16 is exercised by the kernel today (state_t is selected
-# separately via ``STATE_DTYPE_BF16``).
+# Only bf16 is exercised by the kernel today (state_t is selected separately via ``STATE_DTYPE_BF16``).
 _TORCH_DTYPE = {
     "torch.bfloat16": "bfloat16",
     "torch.float16": "float16",
@@ -144,9 +142,8 @@ def parse_csv(csv_path: str) -> list[dict[str, Any]]:
                     "save_vn": _parse_bool(row.get("save_vn") or "True"),
                     "is_varlen": _parse_bool(row.get("is_varlen") or "False"),
                     "wu_contig": _parse_bool(row.get("wu_contig") or "True"),
-                    # state dtype is not tracked in the tuned csv yet; default
-                    # f32 here, then main() unconditionally fans out into a
-                    # bf16 twin so both runtime paths are pre-compiled.
+                    # state dtype is not tracked in the tuned csv yet; default f32 here, then main() unconditionally
+                    # fans out into a bf16 twin so both runtime paths are pre-compiled.
                     "state_bf16": False,
                 }
             except (KeyError, ValueError) as e:
@@ -205,11 +202,9 @@ def _compile_chunk_gdn_h_to_cache(
     torch_dtype = _torch_dtype_for_kernel(dtype)
     state_dtype = torch.bfloat16 if state_bf16 else torch.float32
 
-    # Pick a representative T_flat / N for the dummy tensors. These only
-    # influence the host launch shape, not the compiled artifact, so any
-    # value consistent with BT divisibility works. ``is_varlen`` flips
-    # the kernel's cu_seqlens read path at runtime, but the AOT dummy
-    # tensor shape is identical in both modes, so we use a single T.
+    # Pick a representative T_flat / N for the dummy tensors. These only influence the host launch shape, not the
+    # compiled artifact, so any value consistent with BT divisibility works. ``is_varlen`` flips the kernel's cu_seqlens
+    # read path at runtime, but the AOT dummy tensor shape is identical in both modes, so we use a single T.
     T_flat = BT
     N = 1
     B = N
@@ -224,9 +219,8 @@ def _compile_chunk_gdn_h_to_cache(
     h = torch.empty((B, max(T_flat // BT, 1), H, V, K), device=dev, dtype=torch_dtype)
     h0 = torch.empty((N, H, V, K), device=dev, dtype=state_dtype)
     ht = torch.empty((N, H, V, K), device=dev, dtype=state_dtype)
-    # Variable-length book-keeping tensors. FlyDSL JIT does not accept
-    # ``None`` for tensor slots, so allocate small int32 buffers even when
-    # the kernel branch is disabled.
+    # Variable-length book-keeping tensors. FlyDSL JIT does not accept ``None`` for tensor slots, so allocate small
+    # int32 buffers even when the kernel branch is disabled.
     cu_seqlens = torch.zeros((N + 1,), device=dev, dtype=torch.int32)
     chunk_offsets = torch.zeros((N + 1,), device=dev, dtype=torch.int32)
 
@@ -387,15 +381,13 @@ def main():
 
     all_jobs = collect_aot_jobs(csv_paths, parse_csv)
 
-    # ``--target-arch`` rewrites the ``arch`` field of every job, then we
-    # dedupe again because two csv rows that differed only in arch
-    # collapse to the same compile after the override.
+    # ``--target-arch`` rewrites the ``arch`` field of every job, then we dedupe again because two csv rows that
+    # differed only in arch collapse to the same compile after the override.
     if args.target_arch and all_jobs:
         all_jobs = dedupe_jobs([dict(j, arch=args.target_arch) for j in all_jobs])
 
-    # Fan out into both f32-state and bf16-state variants so neither
-    # runtime path pays a JIT cost on first call. ``dedupe_jobs`` drops
-    # any pre-existing dup.
+    # Fan out into both f32-state and bf16-state variants so neither runtime path pays a JIT cost on first call.
+    # ``dedupe_jobs`` drops any pre-existing dup.
     if all_jobs:
         all_jobs = dedupe_jobs(all_jobs + [dict(j, state_bf16=True) for j in all_jobs])
 

aiter/aot/flydsl/common.py

@@ -157,8 +157,7 @@ def _compile_one_config_for(kind: OpKind) -> Callable[..., dict[str, Any]]:
     elif kind is OpKind.CHUNK_GDN_H:
         from .chunk_gdn_h import compile_one_config
     elif kind is OpKind.GROUPED_MOE:
-        # grouped_moe AOT not wired up yet (no jobs are ever collected); keep a
-        # trivial stub so the dispatch is total.
+        # grouped_moe AOT not wired up yet (no jobs are ever collected); keep a trivial stub so the dispatch is total.
         return lambda **_kw: {}
     else:
         raise ValueError(f"unknown FlyDSL AOT kind: {kind!r}")
@@ -308,9 +307,8 @@ def reap(proc: Any) -> None:
         idx, _ = running.pop(proc)
         out_path = os.path.join(result_dir, f"k{idx}.json")
         if proc.exitcode != 0:
-            # Worker died abnormally (OOM-kill -9, segfault -11, ...) before
-            # writing its result. Transient -> retry (terminal once retries run
-            # out, leaving results[idx]=None).
+            # Worker died abnormally (OOM-kill -9, segfault -11, ...) before writing its result. Transient -> retry
+            # (terminal once retries run out, leaving results[idx]=None).
             retry_or_drop(idx, f"worker crashed (exitcode={proc.exitcode})")
             return
         # Clean exit (exitcode 0): deterministic, never retried.
@@ -358,8 +356,8 @@ def reap(proc: Any) -> None:
 
             launch()  # refill freed slots (including any just-requeued retries)
     finally:
-        # Kill any survivors (e.g. on an unexpected exception) so we never leave
-        # orphaned compilers blocking the caller's exit.
+        # Kill any survivors (e.g. on an unexpected exception) so we never leave orphaned compilers blocking the
+        # caller's exit.
         for proc in list(running):
             try:
                 if proc.is_alive():
@@ -421,9 +419,8 @@ def run_aot(cache_dir: str) -> None:
 
     max_workers = get_max_workers(len(all_jobs))
 
-    # Per-child result files live here -- recreated fresh so stale results
-    # from a previous (e.g. crashed) build can never be mistaken for this
-    # run's output.
+    # Per-child result files live here -- recreated fresh so stale results from a previous (e.g. crashed) build can
+    # never be mistaken for this run's output.
     result_dir = os.path.join(cache_dir, ".aot_results")
     shutil.rmtree(result_dir, ignore_errors=True)
     os.makedirs(result_dir, exist_ok=True)

aiter/aot/flydsl/gemm.py

@@ -278,8 +278,7 @@ def _compile_hgemm_to_cache(
         c_to_lds=c_to_lds,
         has_bias=has_bias,
     )
-    # FlyDSL JIT does not accept None for tensor slots; pass real buffers for
-    # optional bias and split-K sync tensors.
+    # FlyDSL JIT does not accept None for tensor slots; pass real buffers for optional bias and split-K sync tensors.
     launch_bias = bias if has_bias else b
     _compile_executable_to_cache(
         exe,
@@ -440,8 +439,8 @@ def main():
 
     total_t0 = time.time()
 
-    # HGEMM and preshuffle kernels are independent compiles, so they share
-    # one pool for maximum fan-out instead of two serial passes.
+    # HGEMM and preshuffle kernels are independent compiles, so they share one pool for maximum fan-out instead of two
+    # serial passes.
     print(f"\n--- Compiling {len(all_jobs)} kernels (hgemm + preshuffle) ---")
     results = run_jobs_parallel(compile_one_config, hgemm_jobs + preshuffle_jobs)
 

aiter/aot/flydsl/grouped_moe.py

@@ -92,17 +92,14 @@ def _as_int(value, default: int | None = None) -> int | None:
 
 
 def _scheduler_variants(row, base_job):
-    # Production dispatch (grouped_moe_gfx1250._maybe_grouped_gfx1250_a8w4_moe)
-    # hardcodes grouped_persistent_m=False and expert_sched_mode=False; the only
-    # runtime axis is dense vs DeepGEMM contiguous-M (auto-enabled for large token
-    # counts). Mirror exactly that set so AOT never compiles GEMM variants the
-    # runtime cannot launch.
+    # Production dispatch (grouped_moe_gfx1250._maybe_grouped_gfx1250_a8w4_moe) hardcodes grouped_persistent_m=False and
+    # expert_sched_mode=False; the only runtime axis is dense vs DeepGEMM contiguous-M (auto-enabled for large token
+    # counts). Mirror exactly that set so AOT never compiles GEMM variants the runtime cannot launch.
     explicit_contiguous = _as_bool(row.get("grouped_contiguous_m"), False)
     contiguous_modes = [True] if explicit_contiguous else [False, True]
     warp_tile_m = base_job["tile_m"] // base_job["m_warp"]
-    # Contiguous-M sizes max_m as max(cfg.max_m, token_num*topk) (default 0 when
-    # the CSV omits max_m); dense uses cfg.max_m (default token_num). max_m is
-    # baked into the GEMM kernel_tag, so AOT must derive it per-mode exactly as
+    # Contiguous-M sizes max_m as max(cfg.max_m, token_num*topk) (default 0 when the CSV omits max_m); dense uses
+    # cfg.max_m (default token_num). max_m is baked into the GEMM kernel_tag, so AOT must derive it per-mode exactly as
     # grouped_moe_gfx1250 does or the precompiled kernel never gets hit.
     contiguous_max_m = _align_max_m(
         max(_as_int(row.get("max_m"), 0), base_job["token_num"] * base_job["topk"]),
@@ -225,8 +222,8 @@ def _scale_pre_numel(out_e, out_m, feat_dim):
         return out_e * (out_m // wmma_rep) * ((feat_dim // 32) * wmma_rep)
 
     def _route_ksplit(feat_dim, source_topk, out_e, out_m, grouped_in_numel):
-        # build_moe_fused_quant_preshuffle_route_ksplit_module; runtime never
-        # sets remap_rows on the grouped MoE fast path (row_starts stays None).
+        # build_moe_fused_quant_preshuffle_route_ksplit_module; runtime never sets remap_rows on the grouped MoE fast
+        # path (row_starts stays None).
         launch = build_moe_fused_quant_preshuffle_route_ksplit_module(
             feat_dim=feat_dim,
             wmma_rep=wmma_rep,
@@ -285,8 +282,8 @@ def _topids_to_rows():
 
     # --- Stage1 activation prep (a1): fused route + MX-quant + scatter ---
     if contiguous:
-        # DeepGEMM contiguous-M: topids_to_rows -> contiguous psum+remap ->
-        # route-indexed quant+preshuffle into a single contiguous (E=1) buffer.
+        # DeepGEMM contiguous-M: topids_to_rows -> contiguous psum+remap -> route-indexed quant+preshuffle into a single
+        # contiguous (E=1) buffer.
         ub = int(token_num) * int(topk) + int(E) * (int(tile_m) - 1)
         contiguous_m = max(int(tile_m), _align_up(ub, int(tile_m)))
 
@@ -360,9 +357,8 @@ def _topids_to_rows():
         a2_out_e, a2_out_m = E, max_m
 
     # --- Stage2 activation prep (a2): fused grouped quant + preshuffle ---
-    # Runtime passes topids_to_rows whenever route rows fit the output capacity
-    # (almost always), taking the route-ksplit path with source_topk=0; the
-    # plain skip-padding path is the rare capacity-overflow fallback.
+    # Runtime passes topids_to_rows whenever route rows fit the output capacity (almost always), taking the route-ksplit
+    # path with source_topk=0; the plain skip-padding path is the rare capacity-overflow fallback.
     capacity_rows = a2_out_e * a2_out_m
     if numel < capacity_rows:
         _route_ksplit(
@@ -373,8 +369,7 @@ def _topids_to_rows():
             grouped_in_numel=a2_out_e * a2_out_m * inter_dim,
         )
     else:
-        # masked_m is None on the contiguous path -> skip_padding=False;
-        # passed on the dense path -> skip_padding=True.
+        # masked_m is None on the contiguous path -> skip_padding=False; passed on the dense path -> skip_padding=True.
         _plain_preshuffle(
             feat_dim=inter_dim,
             out_e=a2_out_e,
@@ -383,12 +378,9 @@ def _topids_to_rows():
         )
 
     # --- Epilogue: token-order gather-reduce (bf16/f16 fast path only) ---
-    # split_k mirrors stage2's split_k2: when split_k2 > 1 the GEMM emits an
-    # unreduced (split_k, E, max_m, model_dim) tensor and gather-reduce folds the
-    # split slices itself, so its kernel identity depends on split_k. Hardcoding
-    # 1 here makes the token=1 / split_k2>1 CSV rows miss the AOT cache at
-    # inference. vec width is token-count dependent at runtime; cover both so
-    # run-only coverage holds across inference batch sizes, not just the CSV token.
+    # split_k mirrors stage2's split_k2: split_k2>1 emits an unreduced (split_k, E, max_m, model_dim) tensor
+    # gather-reduce folds itself, so its kernel identity depends on split_k (don't hardcode 1). vec width is token-count
+    # dependent; cover both (2, 4) to span inference batch sizes.
     split_k = job["split_k2"]
     for vec in (2, 4):
         gather_reduce = build_moe_gather_reduce_module(
@@ -420,9 +412,8 @@ def compile_one_config(**job):
 
     aot_arch = job.pop("gfx", "") or GROUPED_MOE_AOT_ARCH_DEFAULT
     shape_str = (
-        # Use .get() so a missing key can't raise here, outside the try below:
-        # an escaping exception would crash the worker (exitcode != 0), which the
-        # AOT pool misreads as a transient failure and retries -> potential deadlock.
+        # Use .get() so a missing key can't raise here, outside the try below: an escaping exception would crash the
+        # worker (exitcode != 0), which the AOT pool misreads as a transient failure and retries -> potential deadlock.
         f"{job.get('kernel_name', 'grouped_gemm')}  "
         f"model_dim={job.get('model_dim')} inter_dim={job.get('inter_dim')} "
         f"E={job.get('experts')} topk={job.get('topk')} "
@@ -550,10 +541,9 @@ def compile_one_config(**job):
                 stream=0,
                 _m_tile_map=contiguous_layout,
             )
-            # Bias-epilogue variant: runtime calls stage1(..., bias=...) when the model
-            # carries per-expert bias (e.g. gpt-oss), which triggers a distinct compiled
-            # kernel (gemm1_bias_* / finalize_act_bias). Precompile it alongside the
-            # bias-free kernel so neither path JITs at first inference.
+            # Bias-epilogue variant: runtime calls stage1(..., bias=...) when the model carries per-expert bias (e.g.
+            # gpt-oss), which triggers a distinct compiled kernel (gemm1_bias_* / finalize_act_bias). Precompile it
+            # alongside the bias-free kernel so neither path JITs at first inference.
             bias1 = torch.empty((job["experts"], 2 * job["inter_dim"]), dtype=dtype)
             exe1(
                 y1,
@@ -602,11 +592,9 @@ def compile_one_config(**job):
                 _m_tile_map=contiguous_layout,
                 bias=bias2,
             )
-            # Non-GEMM auxiliary kernels the run-only fast path launches around
-            # the GEMMs (fused route+quant+scatter, grouped quant+preshuffle,
-            # contiguous prefix-sum(+remap), gather-reduce). These were fused in
-            # gfx1250_moe_splitk_fused, so AOT mirrors the new wrappers, not the
-            # legacy route_maps/scatter_copy kernels.
+            # Non-GEMM auxiliary kernels the run-only fast path launches around the GEMMs (fused route+quant+scatter,
+            # grouped quant+preshuffle, contiguous prefix-sum(+remap), gather-reduce). These were fused in
+            # gfx1250_moe_splitk_fused, so AOT mirrors the new wrappers, not the legacy route_maps/scatter_copy kernels.
             _compile_grouped_moe_aux_kernels(
                 job,
                 dtype=dtype,
@@ -618,10 +606,9 @@ def compile_one_config(**job):
         print(f"  [OK] compile  {elapsed:6.1f}s  {shape_str}  arch={aot_arch}")
         return {**job, "compile_time": elapsed, "compile_arch": aot_arch}
     except Exception as e:
-        # Catch everything and return cleanly with compile_time=None: the AOT pool
-        # keys off exitcode 0 + compile_time=None to mark "produced no kernel" and
-        # NOT retry it. An escaping exception crashes the worker (exitcode != 0),
-        # which the pool misreads as a transient failure and retries -> deadlock.
+        # Catch everything and return cleanly with compile_time=None: the AOT pool keys off exitcode 0 +
+        # compile_time=None to mark "produced no kernel" and NOT retry it. An escaping exception crashes the worker
+        # (exitcode != 0), which the pool misreads as a transient failure and retries -> deadlock.
         print(f"  [FAIL] compile  {shape_str}  arch={aot_arch}: {e}")
         traceback.print_exc()
         return {**job, "compile_time": None, "compile_arch": aot_arch}