Skip to content

Conversation

oraluben
Copy link
Contributor

@oraluben oraluben commented Oct 9, 2025

and add FA102 for missing from __future__ import annotations. Prevent issues like #959.

Or use UP035 (with FA102) only if this is too wide.

This should be sufficient to keep py38 compatibility for a long time.

Summary by CodeRabbit

  • New Features

    • Enhanced automatic target detection and gating for advanced CUDA features.
    • Added device metadata fields (platform, compute_capability).
  • Bug Fixes

    • More robust buffer clearing and device attribute fetching with safer fallbacks.
  • Refactor

    • Modernized type hints across the API (PEP 585/604), updated optional/union usage, and streamlined class bases.
    • Standardized file opening defaults and strengthened adapter/profiler signatures.
    • Minor math/logic cleanups without changing core algorithms.
  • Style

    • Replaced legacy string formatting with f-strings across examples and core modules.
  • Chores

    • Updated lint configuration (ruff) and exclusions.

Copy link

github-actions bot commented Oct 9, 2025

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run bash format.sh in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work!

🚀

Copy link
Contributor

coderabbitai bot commented Oct 9, 2025

Walkthrough

This PR modernizes type hints across the codebase to PEP 604/585 syntax, switches many string.format calls to f-strings, and standardizes file open calls to default text mode. Several public APIs have updated signatures. Notable functional tweaks include target auto-detection, phase gating logic, enhanced fill/clear behavior, and CUDA driver error handling.

Changes

Cohort / File(s) Summary
Docs config
docs/conf.py
Minor f-string and open() default-mode changes; behavior unchanged.
Examples: formatting only
examples/* (multiple attention/flash/dequantize/hadamard/norm/linear_attention...)
Switched print formatting to f-strings; some comprehensions to generators; no logic change.
Examples: typing updates
examples/attention_sink/*, examples/cast/*, examples/flash_decoding/*, examples/deepseek_*/*, examples/grouped_gemm/*
Added future annotations; Optional/Tuple/Dict→union/built-in generics; minor signature updates; behavior intact.
BitNet 1.58b
examples/bitnet-1.58b/*
Removed coding lines; adopted modern typing; open() default mode; numerous public signatures updated to unions/built-in generics; no algorithm changes.
Blocksparse examples
examples/blocksparse_*/*
super() modernization and f-strings in diagnostics; no logic change.
Tests
testing/python/*
Minor parenthesis cleanup and Optional→int
Tooling/config
pyproject.toml, setup.py
Ruff config tweaks (enable UP/FA, ignore SIM115, exclusions); type hint modernizations; open() default mode.
Autotuner core
tilelang/autotuner/*
Future annotations; Optional/Union/List→modern types; public class/func signatures updated; behavior preserved.
Kernel cache
tilelang/cache/*
Future annotations; signature types modernized; open() default mode; no functional change.
JIT API (public surface)
tilelang/jit/__init__.py, .../kernel.py, .../adapter/*, .../adapter/utils.py, .../adapter/wrapper.py, .../adapter/libgen.py, .../adapter/cython/*, .../adapter/nvrtc/*, .../adapter/ctypes/*, .../adapter/torch/metal.py
Broad migration to PEP 604/585 types; many public signatures updated (targets, out_idx, pass_configs, compile_flags); file opening defaults; no core flow changes.
Engine
tilelang/engine/lower.py, .../param.py, .../phase.py, .../callback.py
Modernized types; phase gating functions updated with target/TMA checks and env flags; public signatures adjusted.
Language: core APIs
tilelang/language/__init__.py, atomic.py, builtin.py, copy.py, customize.py, fill.py, frame.py, gemm.py, kernel.py, logical.py, parallel.py, persistent.py, pipeline.py, reduce.py, tir/*, utils.py, parser/operation.py, experimental/gemm_sp.py, warpgroup.py
Widespread type updates; some API signatures widened; fill.clear enhanced to handle Var-bound regions; kernel/frame accessors return list[...] types; minor assertion/message f-strings.
Carver: arch
tilelang/carver/arch/__init__.py, arch_base.py, cuda.py, cdna.py, metal.py, driver/cuda_driver.py
Type modernizations; added platform/compute_capability fields; CDNA/CUDA constructors hint updates; CUDA driver now returns None on attribute errors; behavior otherwise same.
Carver: roller/analysis/template
tilelang/carver/analysis.py, roller/*, template/*, utils.py, matmul_analysis.py, common_schedules.py, shape_inference/*, rasterization.py
Modern typing across public classes/methods; removal of object base; cached decorator change (@lru_cache); signatures updated; logic preserved.
Intrinsics
tilelang/intrinsics/*
Future annotations; Optional/Tuple→modern types; minor f-string messages; public constructor/signature updates; behavior unchanged.
Primitives GEMM
tilelang/primitives/gemm/*, tilelang/tileop/gemm/gemm_base.py
Type updates to unions; GemmBase changed from dataclass to plain class (dataclass methods removed).
Quantization
tilelang/quantize/*
Return types to dict/tuple; f-string messages; minor parentheses cleanup; behavior same.
Profiler
tilelang/profiler/*
Public API signatures use modern types; helper returns updated; behavior intact.
Contrib/Tools
tilelang/contrib/*, tilelang/tools/Analyzer.py
Type modernizations; minor f-strings; unchanged logic.
Utils/Target
tilelang/utils/*
Type modernizations; determine_target now resolves "auto" more explicitly; other helpers updated.
Versioning
tilelang/version.py
Future annotations; return type str

Sequence Diagram(s)

sequenceDiagram
  actor User
  participant API as jit.compile/cached
  participant Target as determine_target
  participant Lower as engine.lower
  participant Adapter as JIT Adapter
  participant Kernel as JITKernel

  User->>API: call with target="auto"
  API->>Target: resolve target
  alt auto-detect
    Target-->>API: concrete Target (cuda/hip/metal) or error
  else explicit
    Target-->>API: validated Target/str
  end
  API->>Lower: lower(func_or_mod, target, target_host, pass_configs)
  Lower-->>API: CompiledArtifact
  API->>Adapter: create backend adapter (dlpack/ctypes/cython/nvrtc)
  Adapter-->>Kernel: instantiate runnable kernel
  Kernel-->>User: callable kernel/reference pair
Loading
sequenceDiagram
  participant TL as tl.fill/clear
  participant Utils as utils.get_buffer_region_from_load
  participant FFI as _ffi_api

  TL->>TL: clear(buffer)
  alt buffer is Buffer
    TL->>FFI: fill(buffer, 0)
  else buffer is Var(let-bound)
    TL->>TL: inspect let value
    alt BufferRegion
      TL->>FFI: fill(region, 0)
    else BufferLoad
      TL->>Utils: get_buffer_region_from_load(load)
      alt valid region
        TL->>FFI: fill(region, 0)
      else invalid
        TL-->>TL: raise ValueError
      end
    else other
      TL-->>TL: raise ValueError
    end
  end
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

Suggested reviewers

  • LeiWang1999

Poem

Hop-hop I tidy types, through fields and flows,
Swapping Optional’s burrows for unions’ new clothes.
F-strings twinkle, open() breathes light—
Targets auto-find their starry night.
Kernels compiled, caches warm—
A rabbit reviews in typist form. 🐇✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)
Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%. You can run @coderabbitai generate docstrings to improve docstring coverage.
✅ Passed checks (2 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check ✅ Passed The title “[Lint] Enable pyupgrade linter in ruff” accurately and succinctly captures the principal change of activating pyupgrade-related rules in the ruff configuration, reflecting the updates in pyproject.toml. It is clearly related to the core intent of the PR and avoids unnecessary details. A team member glancing at the commit history will immediately understand that this PR focuses on lint settings for pyupgrade enforcement.
✨ Finishing touches
  • 📝 Generate docstrings
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 3d0803a and 013834f.

📒 Files selected for processing (3)
  • pyproject.toml (2 hunks)
  • setup.py (9 hunks)
  • tilelang/jit/adapter/libgen.py (5 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
  • pyproject.toml
🧰 Additional context used
🧬 Code graph analysis (2)
setup.py (1)
tilelang/jit/adapter/cython/adapter.py (1)
  • get_cython_compiler (34-66)
tilelang/jit/adapter/libgen.py (3)
tilelang/jit/adapter/ctypes/adapter.py (3)
  • srcpath (270-272)
  • libpath (275-277)
  • lib_code (280-282)
tilelang/jit/adapter/cython/adapter.py (3)
  • srcpath (502-504)
  • libpath (507-509)
  • lib_code (512-514)
tilelang/jit/adapter/wrapper.py (2)
  • assign_pass_configs (1105-1106)
  • host_func (626-636)
🪛 Ruff (0.13.3)
setup.py

312-312: Unused noqa directive (non-enabled: SIM115)

Remove unused noqa directive

(RUF100)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
  • GitHub Check: build-test-metal
  • GitHub Check: format-check
🔇 Additional comments (8)
setup.py (3)

1-1: LGTM! Correct placement of future annotations import.

This enables postponed evaluation of type annotations, allowing modern type hint syntax (e.g., list[str], str | None) to work with Python 3.8 as intended by the PR objectives.


124-124: LGTM! Type hints modernized correctly.

The updates from List[str] to list[str] (line 124) and Optional[str] to str | None (line 246) align with PEP 585 and PEP 604. These are safe for Python 3.8 with the __future__ import and consistent with changes across the codebase.

Also applies to: 246-246


140-140: LGTM! Simplified file open calls.

Removing the explicit "r" mode is correct since text mode is the default for open(). These changes make the code more idiomatic without any functional impact.

Also applies to: 171-171, 711-711, 723-723, 740-740

tilelang/jit/adapter/libgen.py (5)

1-1: LGTM! Future annotations import enables modern type hints.

The addition of from __future__ import annotations is essential for using built-in generics (dict, list) and PEP 604 unions (|) in type hints while maintaining Python 3.8 compatibility through postponed annotation evaluation.


33-33: LGTM! Modernized class signature.

Removing explicit object inheritance is correct for Python 3+, as all classes implicitly inherit from object.


34-38: LGTM! Type annotations correctly modernized.

The class attribute type annotations have been properly updated to use:

  • PEP 604 union syntax (str | None instead of Optional[str])
  • Built-in generics (dict, list instead of Dict, List)

These changes are compatible with Python 3.8+ when using from __future__ import annotations.


44-56: LGTM! Method signatures correctly updated.

The method parameter type annotations have been properly modernized to match the class attribute types using PEP 604 unions and built-in generics.


189-210: LGTM! PyLibraryGenerator changes are consistent.

The type annotation updates for PyLibraryGenerator are consistent with the parent class LibraryGenerator, correctly using modern type hint syntax.


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@oraluben oraluben marked this pull request as ready for review October 9, 2025 23:56
Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
tilelang/jit/adapter/utils.py (1)

78-81: Guard attrs["global_symbol"] access on PrimFunc
Accessing func_or_mod.attrs["global_symbol"] may KeyError if the attribute isn’t set (e.g. PrimFunc created in parser). Use attrs.get("global_symbol") with a clear fallback or assert its presence before indexing.

🧹 Nitpick comments (2)
tilelang/carver/roller/shape_inference/tir.py (1)

354-354: LGTM: Modernized string formatting.

The f-string is clearer and more readable than the old % formatting style.

For completeness, the static analysis tool suggests considering TypeError for invalid type errors and potentially creating a custom exception class, but these are minor style improvements that can be deferred.

tilelang/carver/roller/node.py (1)

304-304: LGTM: Simplified decorator syntax.

The simplified @functools.lru_cache (without parentheses) is valid and cleaner in Python 3.8+.

Note: Static analysis warns that using lru_cache on instance methods can lead to memory leaks because the cache holds references to self, preventing garbage collection. This is an existing pattern in the codebase and not introduced by this change, but consider whether these methods truly need caching on instance methods or if the cache should be cleared when instances are no longer needed.

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a13cde2 and 0973611.

📒 Files selected for processing (72)
  • docs/conf.py (1 hunks)
  • examples/amd/example_amd_flash_attn_bwd.py (2 hunks)
  • examples/attention_sink/example_gqa_sink_bwd_bhsd.py (3 hunks)
  • examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py (2 hunks)
  • examples/attention_sink/example_mha_sink_bwd_bhsd.py (3 hunks)
  • examples/attention_sink/example_mha_sink_fwd_bhsd.py (2 hunks)
  • examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py (2 hunks)
  • examples/bitnet-1.58b/configuration_bitnet.py (0 hunks)
  • examples/bitnet-1.58b/eval_ppl.py (1 hunks)
  • examples/bitnet-1.58b/maint/create_bitblas_ckpt.py (1 hunks)
  • examples/bitnet-1.58b/modeling_bitnet.py (1 hunks)
  • examples/bitnet-1.58b/tokenization_bitnet.py (0 hunks)
  • examples/bitnet-1.58b/utils_quant.py (1 hunks)
  • examples/bitnet-1.58b/vllm_workspace/conftest.py (1 hunks)
  • examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py (2 hunks)
  • examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py (2 hunks)
  • examples/cast/example_group_per_split_token_cast_to_fp8.py (1 hunks)
  • examples/cast/example_per_token_cast_to_fp8.py (2 hunks)
  • examples/deepseek_mla/example_mla_decode_paged.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_bf16_fp4_hopper.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper_tma.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_fp4_hopper.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_groupedgemm_bf16_mxfp4_hopper.py (1 hunks)
  • examples/flash_attention/example_gqa_bwd.py (2 hunks)
  • examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (2 hunks)
  • examples/flash_attention/example_gqa_fwd_bshd.py (1 hunks)
  • examples/flash_attention/example_gqa_fwd_bshd_wgmma_pipelined.py (1 hunks)
  • examples/flash_attention/example_mha_bwd.py (2 hunks)
  • examples/flash_attention/example_mha_bwd_bhsd.py (2 hunks)
  • examples/flash_attention/example_mha_bwd_wgmma_pipelined.py (2 hunks)
  • examples/flash_attention/example_mha_fwd_bhsd.py (1 hunks)
  • examples/flash_attention/example_mha_fwd_bhsd_wgmma_pipelined.py (1 hunks)
  • examples/flash_attention/example_mha_fwd_bshd.py (1 hunks)
  • examples/flash_attention/example_mha_fwd_bshd_wgmma_pipelined.py (1 hunks)
  • examples/flash_decoding/example_gqa_decode.py (1 hunks)
  • examples/flash_decoding/example_mha_inference.py (1 hunks)
  • examples/grouped_gemm/example_grouped_gemm_bwd.py (1 hunks)
  • examples/hadamard_transform/example_hadamard.py (1 hunks)
  • examples/linear_attention/example_mamba_chunk_scan.py (1 hunks)
  • examples/linear_attention/example_mamba_chunk_state.py (2 hunks)
  • examples/minference/example_vertical_slash_sparse_attn.py (1 hunks)
  • examples/norm/rms_norm.py (1 hunks)
  • pyproject.toml (1 hunks)
  • setup.py (6 hunks)
  • testing/python/kernel/test_tilelang_kernel_gemm.py (1 hunks)
  • testing/python/kernel/test_tilelang_kernel_gemm_simt.py (1 hunks)
  • testing/python/language/test_tilelang_language_pipeline.py (1 hunks)
  • tilelang/autotuner/param.py (3 hunks)
  • tilelang/cache/kernel_cache.py (1 hunks)
  • tilelang/carver/arch/cuda.py (1 hunks)
  • tilelang/carver/arch/metal.py (1 hunks)
  • tilelang/carver/roller/bestfit.py (1 hunks)
  • tilelang/carver/roller/hint.py (1 hunks)
  • tilelang/carver/roller/node.py (4 hunks)
  • tilelang/carver/roller/rasterization.py (1 hunks)
  • tilelang/carver/roller/shape_inference/common.py (2 hunks)
  • tilelang/carver/roller/shape_inference/tir.py (3 hunks)
  • tilelang/contrib/hipcc.py (1 hunks)
  • tilelang/intrinsics/mfma_macro_generator.py (1 hunks)
  • tilelang/intrinsics/mma_macro_generator.py (2 hunks)
  • tilelang/intrinsics/wgmma_macro_generator.py (1 hunks)
  • tilelang/jit/adapter/cython/adapter.py (3 hunks)
  • tilelang/jit/adapter/libgen.py (1 hunks)
  • tilelang/jit/adapter/utils.py (3 hunks)
  • tilelang/jit/adapter/wrapper.py (6 hunks)
  • tilelang/jit/kernel.py (1 hunks)
  • tilelang/language/proxy.py (4 hunks)
  • tilelang/quantize/lop3.py (1 hunks)
  • tilelang/quantize/quantization.py (2 hunks)
  • tilelang/tileop/gemm/gemm_base.py (1 hunks)
  • tilelang/version.py (1 hunks)
💤 Files with no reviewable changes (2)
  • examples/bitnet-1.58b/tokenization_bitnet.py
  • examples/bitnet-1.58b/configuration_bitnet.py
🧰 Additional context used
🧬 Code graph analysis (25)
examples/minference/example_vertical_slash_sparse_attn.py (1)
tilelang/language/builtin.py (1)
  • mbarrier_wait_parity (172-219)
examples/flash_attention/example_mha_fwd_bshd_wgmma_pipelined.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_mha_fwd_bshd.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_mha_fwd_bhsd_wgmma_pipelined.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
tilelang/carver/roller/shape_inference/common.py (1)
tilelang/carver/roller/shape_inference/tir.py (2)
  • Statement (7-43)
  • InputShapeInference (169-318)
examples/flash_decoding/example_mha_inference.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
tilelang/intrinsics/wgmma_macro_generator.py (1)
tilelang/utils/language.py (1)
  • is_fragment (68-78)
examples/flash_attention/example_gqa_fwd_bshd_wgmma_pipelined.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/attention_sink/example_mha_sink_bwd_bhsd.py (2)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (2)
  • maybe_contiguous (356-359)
  • tl_bwd (484-485)
examples/linear_attention/example_mamba_chunk_state.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_mha_bwd_wgmma_pipelined.py (4)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (1)
  • maybe_contiguous (356-359)
examples/flash_attention/example_gqa_bwd.py (1)
  • maybe_contiguous (377-380)
examples/grouped_gemm/example_grouped_gemm_bwd.py (1)
  • maybe_contiguous (134-137)
examples/dequantize_gemm/example_dequant_gemm_fp4_hopper.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_decoding/example_gqa_decode.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/linear_attention/example_mamba_chunk_scan.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
tilelang/intrinsics/mma_macro_generator.py (2)
tilelang/intrinsics/wgmma_macro_generator.py (1)
  • TensorCoreIntrinEmitter (64-519)
tilelang/utils/language.py (1)
  • is_fragment (68-78)
examples/flash_attention/example_mha_bwd_bhsd.py (1)
examples/flash_attention/example_mha_bwd.py (1)
  • maybe_contiguous (258-261)
examples/flash_attention/example_mha_fwd_bhsd.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/grouped_gemm/example_grouped_gemm_bwd.py (1)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (3)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/attention_sink/example_mha_sink_bwd_bhsd.py (1)
  • maybe_contiguous (364-367)
examples/flash_attention/example_gqa_bwd.py (1)
  • maybe_contiguous (377-380)
examples/flash_attention/example_mha_bwd.py (9)
examples/amd/example_amd_flash_attn_bwd.py (2)
  • maybe_contiguous (242-245)
  • run1 (340-341)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (1)
  • maybe_contiguous (356-359)
examples/attention_sink/example_mha_sink_bwd_bhsd.py (1)
  • maybe_contiguous (364-367)
examples/flash_attention/example_gqa_bwd.py (2)
  • maybe_contiguous (377-380)
  • run1 (514-515)
examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (2)
  • maybe_contiguous (399-402)
  • run1 (536-537)
examples/flash_attention/example_mha_bwd_bhsd.py (2)
  • maybe_contiguous (259-262)
  • run1 (336-337)
examples/flash_attention/example_mha_bwd_wgmma_pipelined.py (2)
  • maybe_contiguous (268-271)
  • run1 (343-344)
examples/grouped_gemm/example_grouped_gemm_bwd.py (1)
  • maybe_contiguous (134-137)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/amd/example_amd_flash_attn_bwd.py (1)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (1)
  • maybe_contiguous (356-359)
examples/flash_attention/example_gqa_fwd_bshd.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
tilelang/jit/adapter/utils.py (1)
tilelang/language/ast/ir.py (1)
  • target (1682-1713)
examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (1)
examples/flash_attention/example_gqa_bwd.py (2)
  • maybe_contiguous (377-380)
  • run1 (514-515)
examples/flash_attention/example_gqa_bwd.py (1)
examples/amd/example_amd_flash_attn_bwd.py (2)
  • maybe_contiguous (242-245)
  • run1 (340-341)
🪛 GitHub Actions: CI
examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py

[error] 1-1: Reformatted by formatter. Please review and stage the changes.

docs/conf.py

[error] 1-1: Reformatted by formatter. Please review and stage the changes.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py

[error] 1-1: Reformatted by formatter. Please review and stage the changes.

🪛 GitHub Actions: CI Test on AMD
examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py

[error] 1-1: Reformatted files. Please review and stage the changes.

docs/conf.py

[error] 1-1: Reformatted files. Please review and stage the changes.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py

[error] 1-1: Reformatted files. Please review and stage the changes.

🪛 GitHub Actions: CI Test on Metal
examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py

[error] 1-1: Reformatted by formatting script. Changes not staged for commit.

docs/conf.py

[error] 1-1: Reformatted by formatting script. Changes not staged for commit.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py

[error] 1-1: Reformatted by formatting script. Changes not staged for commit.

🪛 Ruff (0.13.3)
tilelang/carver/roller/node.py

304-304: Use of functools.lru_cache or functools.cache on methods can lead to memory leaks

(B019)


421-421: Use of functools.lru_cache or functools.cache on methods can lead to memory leaks

(B019)

tilelang/quantize/lop3.py

1189-1189: Avoid specifying long messages outside the exception class

(TRY003)

setup.py

312-312: Unused noqa directive (non-enabled: SIM115)

Remove unused noqa directive

(RUF100)

tilelang/carver/roller/shape_inference/tir.py

354-354: Prefer TypeError exception for invalid type

(TRY004)


354-354: Create your own exception

(TRY002)


354-354: Avoid specifying long messages outside the exception class

(TRY003)

🔇 Additional comments (96)
examples/linear_attention/example_mamba_chunk_state.py (1)

171-175: Print modernisation looks good

Nice switch to f-strings—same formatting semantics, cleaner syntax.

examples/flash_attention/example_mha_fwd_bshd_wgmma_pipelined.py (1)

198-202: Nice f-string modernization.

Formatting stays identical while matching the new lint rule scope. 👍

examples/flash_attention/example_gqa_fwd_bshd.py (1)

251-255: F-string migration looks good.

Formatting precision and computed values are unchanged; the output stays consistent while modernizing the style.

examples/flash_attention/example_mha_fwd_bhsd_wgmma_pipelined.py (1)

211-215: F-string conversion looks good.

Formatting specifiers are preserved, so output remains identical while complying with pyupgrade.

examples/bitnet-1.58b/eval_ppl.py (1)

57-57: LGTM! Clean f-string conversion.

The modernization from legacy string formatting to f-strings improves readability and aligns with Python 3.6+ best practices. This change is consistent with the PR's goal of enabling pyupgrade linter rules.

examples/hadamard_transform/example_hadamard.py (1)

154-154: LGTM! Clean modernization to f-string.

The conversion from .format() to f-string improves readability and aligns with modern Python best practices while maintaining identical functionality.

tilelang/carver/roller/bestfit.py (1)

20-20: LGTM! Modernization to f-string improves readability.

The conversion from str.format to an f-string is correct and aligns with the PR objective to enable pyupgrade linter. F-strings are more readable and typically faster than older formatting methods.

examples/bitnet-1.58b/utils_quant.py (1)

219-219: LGTM! Modernized super() call syntax.

The change from super(BitLinear, self).__init__(*kargs, **kwargs) to super().__init__(*kargs, **kwargs) correctly modernizes the code to use Python 3+ idiomatic syntax. The behavior remains identical.

tilelang/carver/arch/metal.py (1)

1-1: LGTM!

Adding from __future__ import annotations enables postponed evaluation of type annotations (PEP 563), which is a modern Python practice and aligns with the PR's objective to modernize the codebase. This works well with the existing PEP 604 union syntax (Target | str) on line 12.

examples/deepseek_mla/example_mla_decode_paged.py (1)

403-404: LGTM! Clean modernization to f-strings.

The conversion from format() to f-strings is correct and improves readability while maintaining identical output semantics.

examples/flash_decoding/example_gqa_decode.py (1)

476-480: LGTM! Clean modernization to f-strings.

The conversion from .format() to f-strings is correct and follows Python best practices. F-strings are more readable, concise, and performant.

examples/flash_attention/example_gqa_fwd_bshd_wgmma_pipelined.py (2)

225-226: LGTM! Clean conversion to f-strings.

The migration from .format() to f-strings is correct and follows modern Python conventions. The formatting specifiers and output remain identical.


228-229: LGTM! Consistent f-string conversion.

The conversion maintains identical output while improving code readability.

examples/linear_attention/example_mamba_chunk_scan.py (1)

232-236: LGTM! Clean modernization to f-strings.

The conversion from .format() to f-strings improves readability while preserving the exact formatting (.2f). This aligns with the PR objective to enable pyupgrade linting and follows Python 3.6+ best practices.

examples/flash_attention/example_mha_fwd_bhsd.py (1)

206-210: LGTM! Clean modernization to f-strings.

The conversion from .format() to f-strings is correct, maintains identical output formatting (.2f precision), and improves readability. These changes align with the PR objective of enabling pyupgrade linter rules.

examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper.py (1)

537-538: LGTM! Clean f-string conversion.

The formatting is preserved (.2f for 2 decimal places) and the f-string syntax is more readable and Pythonic.

examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper_tma.py (1)

553-554: LGTM! Consistent f-string conversion.

The change maintains identical formatting and improves code readability.

examples/dequantize_gemm/example_dequant_gemm_bf16_fp4_hopper.py (1)

436-437: LGTM! Consistent f-string conversion.

The formatting is identical and the f-string syntax is cleaner and more maintainable.

examples/flash_attention/example_mha_fwd_bshd.py (1)

193-197: LGTM! Clean modernization to f-strings.

The conversion from .format() to f-strings is correct and improves readability. The formatting specifications and variable references are properly preserved.

examples/cast/example_per_token_cast_to_fp8.py (1)

103-117: LGTM! Clean f-string modernization.

The conversion from .format() to f-strings is correct and follows modern Python best practices. The format specifications are preserved, and the output remains identical.

examples/cast/example_group_per_split_token_cast_to_fp8.py (1)

202-205: LGTM! Consistent f-string modernization.

The f-string conversions are correct and consistent with the broader codebase modernization in this PR.

tilelang/intrinsics/mfma_macro_generator.py (1)

14-14: LGTM! Modern Python 3 class declaration.

Removing the explicit object base class aligns with Python 3 best practices, as all classes implicitly inherit from object.

tilelang/intrinsics/wgmma_macro_generator.py (1)

408-408: LGTM! Cleaner assertion message with f-string.

The f-string format improves readability and aligns with modern Python conventions.

tilelang/intrinsics/mma_macro_generator.py (2)

28-28: LGTM! Modern Python 3 class declaration.

Removing the explicit object base class aligns with Python 3 best practices, consistent with the modernization in mfma_macro_generator.py.


521-521: LGTM! Cleaner assertion message with f-string.

The f-string format improves readability and is consistent with the same change in wgmma_macro_generator.py.

examples/amd/example_amd_flash_attn_bwd.py (2)

247-247: LGTM! Generator expression improves efficiency.

The change from list comprehension to generator expression is a good modernization. Since the generator is consumed immediately during unpacking, behavior is unchanged while memory efficiency is slightly improved.


346-350: LGTM! F-string conversions improve readability.

The conversion from .format() to f-strings is correct, with all formatting specifications (.2f) properly preserved. This modernization improves code readability without changing behavior.

examples/flash_attention/example_mha_bwd_wgmma_pipelined.py (2)

273-273: LGTM! Modern generator expression for unpacking.

The change from a list comprehension to a generator expression is a good modernization. The generator is more memory-efficient since the values are immediately consumed during unpacking, and this aligns with pyupgrade recommendations.


349-353: LGTM! F-string conversions are correct.

The conversion to f-strings is a modern Python best practice that improves readability. All formatting specifications (.2f) are correctly preserved, and the logic remains unchanged.

examples/flash_attention/example_mha_bwd_bhsd.py (2)

264-264: LGTM! Generator expression for unpacking is correct.

The change from list comprehension to generator expression is a valid modernization. While the memory benefit is minimal for 5 items, this aligns with pyupgrade's recommendations and works correctly.


342-346: LGTM! F-string conversions are correct.

The print statements have been properly converted from .format() to f-strings with correct formatting specifiers preserved.

tilelang/quantize/quantization.py (2)

226-226: LGTM! Redundant parentheses removed.

The removal of the outer parentheses is a safe cosmetic improvement. The expression remains functionally identical, and the remaining parentheses correctly ensure the bitwise AND operation is evaluated before the shift.


235-235: LGTM! Consistent style improvement.

The redundant outer parentheses have been removed, making the expression cleaner while maintaining the correct evaluation order. This change is consistent with the improvement on line 226.

examples/flash_attention/example_mha_bwd.py (2)

263-263: LGTM! Generator expression modernization.

The change from list comprehension to generator expression for unpacking is a safe, standard Python modernization. Both are functionally equivalent when unpacking, but the generator expression is more memory-efficient.


339-343: LGTM! F-string conversion.

The conversion from .format() to f-strings is a standard Python modernization that improves readability. All format specifiers and expressions are correctly preserved.

examples/flash_attention/example_gqa_bwd.py (2)

382-382: LGTM! Generator expression is more memory-efficient.

The change from list comprehension to generator expression for unpacking is a valid modernization. Both are functionally equivalent, and the generator expression avoids creating an intermediate list.


520-524: LGTM! F-strings improve readability.

The conversion to f-strings modernizes the code and improves readability while correctly preserving the format specifiers for floating-point precision.

examples/grouped_gemm/example_grouped_gemm_bwd.py (1)

139-139: LGTM! Standard pyupgrade optimization.

Replacing the list comprehension with a generator expression is correct and avoids allocating an intermediate list. This is a standard pyupgrade rule (UP015) for immediate unpacking.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py (4)

183-183: LGTM! Modernized super() call.

The simplified super().__init__() syntax is the Python 3 standard and is preferred over the explicit class/self parameters.


357-357: LGTM! F-string conversion.

The f-string syntax is preferred over .format() and improves readability.


362-362: LGTM! F-string conversion.

The f-string syntax with multiple interpolations is more readable than the equivalent .format() call.


1-1: Stage formatting changes.

CI is failing due to unstaged formatter updates. Run ruff --fix (or black . if used) and commit all modified files.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py (3)

196-196: LGTM! Modernized super() call.

The simplified super().__init__() syntax is the Python 3 standard and is preferred over the explicit class/self parameters.


369-369: LGTM! F-string conversion.

The f-string syntax is preferred over .format() and improves readability.


372-372: LGTM! F-string conversion.

The f-string syntax with multiple interpolations is more readable than the equivalent .format() call.

tilelang/contrib/hipcc.py (1)

57-57: F-string update preserves behavior

Switching to the f-string keeps the target path identical while satisfying pyupgrade.

testing/python/language/test_tilelang_language_pipeline.py (1)

106-107: Redundant parentheses removal is safe

Dropping the extra parentheses leaves the casting logic untouched; nice stylistic cleanup.

tilelang/carver/roller/shape_inference/common.py (1)

7-22: Modern class declarations look good

Removing the explicit (object) base aligns with Python 3 style and has no runtime impact.

examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py (1)

3-427: Formatting modernization retains semantics

The added annotations import supports future typing tweaks, and the f-strings mirror the prior formatting exactly.

tilelang/version.py (1)

24-25: Default text-mode open is equivalent

Dropping the explicit 'r' keeps text reading behavior unchanged while satisfying the linter.

examples/bitnet-1.58b/vllm_workspace/conftest.py (1)

37-39: Simplified file open stays readable

Using the default read mode is fine here and keeps the helper neat.

testing/python/kernel/test_tilelang_kernel_gemm.py (1)

98-99: Parentheses trim keeps conversion intact

The tensor bitcast still executes exactly the same, so this satisfies the formatter without behavioral change.

tilelang/carver/roller/rasterization.py (1)

91-93: F-string emission is straightforward

Injecting panel_width via an f-string preserves the generated code verbatim.

pyproject.toml (1)

34-34: LGTM! Pyupgrade linter enabled successfully.

The addition of "UP" and "FA102" to the ruff lint selection enables pyupgrade rules, which modernize Python syntax. This change aligns with the PR objectives and the style updates throughout the codebase.

tilelang/jit/adapter/cython/adapter.py (1)

105-105: LGTM! Simplified file open calls.

Removing the explicit "r" mode is consistent with Python 3 conventions, where text read mode is the default. The behavior remains identical.

Also applies to: 118-118, 135-135

tilelang/autotuner/param.py (1)

249-249: LGTM! Simplified file open calls.

Removing the explicit "r" mode follows Python 3 conventions. The default text read mode preserves the intended behavior.

Also applies to: 318-318, 330-330

setup.py (1)

140-140: LGTM! Simplified file open calls.

Removing the explicit "r" mode follows Python 3 conventions and is consistent with pyupgrade recommendations.

Also applies to: 171-171, 711-711, 723-723, 740-740

tilelang/carver/arch/cuda.py (1)

94-94: LGTM! Removed redundant base class.

In Python 3, all classes implicitly inherit from object, making the explicit base class declaration unnecessary. This modernization aligns with PEP 8 and pyupgrade recommendations.

examples/bitnet-1.58b/modeling_bitnet.py (1)

1515-1515: LGTM! Simplified file open call.

Removing the explicit "r" mode is consistent with Python 3 conventions and the modernization pattern applied throughout the codebase.

examples/bitnet-1.58b/maint/create_bitblas_ckpt.py (1)

88-88: LGTM! Simplified file open call.

The change follows Python 3 conventions and is consistent with the modernization applied across the codebase.

examples/norm/rms_norm.py (1)

74-74: LGTM! Modernized string formatting.

Replacing .format() with f-strings improves readability and performance. This is a recommended Python 3.6+ practice and aligns with pyupgrade recommendations.

Also applies to: 76-76

tilelang/cache/kernel_cache.py (1)

366-368: No behavior change from dropping "r"

open() defaults to text-read mode, so this remains a safe, no-op cleanup.

tilelang/quantize/lop3.py (1)

1189-1193: F-string conversion looks good

Switching to an f-string keeps the message identical while matching the new lint expectations.

docs/conf.py (1)

5-9: Docs tweaks are sound

Both the f-string and the implicit read mode are clean modernizations with no functional impact.

tilelang/carver/roller/hint.py (1)

150-152: Class declaration cleanup is correct

Removing (object) is idiomatic in Python 3 and keeps semantics identical.

tilelang/jit/kernel.py (1)

20-32: Python 3 class style looks good

Dropping the explicit object base keeps behavior the same while modernizing the declaration.

examples/attention_sink/example_mha_sink_fwd_bhsd.py (2)

2-2: LGTM: Future annotations import added.

Enables PEP 563 postponed evaluation of annotations, consistent with modern Python 3 practices.


289-293: LGTM: Format strings modernized to f-strings.

The conversion from .format() to f-strings is correct and improves readability. The numeric formatting and calculations remain unchanged.

examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py (1)

3-3: LGTM: Modern Python formatting applied.

The addition of future annotations and conversion to f-strings are standard modernizations that improve code readability without changing functionality.

Also applies to: 438-446

testing/python/kernel/test_tilelang_kernel_gemm_simt.py (1)

109-109: LGTM: Removed redundant parentheses.

The extra parentheses around the integer expression were unnecessary and have been correctly removed.

examples/attention_sink/example_gqa_sink_bwd_bhsd.py (2)

2-2: LGTM: Modern Python patterns applied.

The future annotations import and generator expression (instead of list comprehension) are appropriate modernizations. The generator expression is safe here since values are immediately unpacked into separate variables.

Also applies to: 361-361


488-492: LGTM: Format strings modernized to f-strings.

The conversion maintains the same output formatting while improving readability.

examples/flash_decoding/example_mha_inference.py (1)

321-325: LGTM: Format strings modernized to f-strings.

The conversion correctly maintains the formatting precision (.2f and .4f) while improving code readability.

examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (2)

404-404: LGTM: Generator expression applied.

The switch from list comprehension to generator expression is more memory efficient and safe here since values are immediately unpacked.


542-546: LGTM: Format strings modernized to f-strings.

The conversion maintains the same output formatting while improving readability.

examples/attention_sink/example_mha_sink_bwd_bhsd.py (3)

2-2: LGTM: Future annotations import added.

Enables PEP 563 forward reference support, consistent with modern Python typing practices.


369-369: LGTM: Generator expression unpacking.

More memory-efficient than creating an intermediate list for unpacking, while maintaining the same functionality.


492-496: LGTM: Modernized to f-strings.

Cleaner and more readable than .format() calls, consistent with modern Python style.

tilelang/carver/roller/shape_inference/tir.py (2)

50-50: LGTM: Removed redundant explicit base class.

In Python 3, all classes implicitly inherit from object, so the explicit base is unnecessary.


79-79: LGTM: Removed redundant explicit base class.

Consistent with modern Python 3 style where object inheritance is implicit.

tilelang/carver/roller/node.py (3)

32-32: LGTM: Removed redundant explicit base class.

Modern Python 3 style where object inheritance is implicit.


93-93: LGTM: Removed redundant explicit base class.

Consistent with Python 3 conventions.


421-421: LGTM: Simplified decorator syntax.

Same modernization as line 304, consistent with Python 3.8+ conventions.

tilelang/jit/adapter/wrapper.py (6)

179-179: LGTM: Removed redundant explicit base class.

Modern Python 3 style.


325-326: LGTM: Modernized to f-strings.

Cleaner and more readable for constructing the dim3 strings.


350-351: LGTM: Modernized to f-strings.

Consistent string formatting improvements for kernel launch and error checking.


895-895: LGTM: Removed redundant explicit base class.

Consistent with Python 3 conventions.


993-993: LGTM: Modernized to f-string.

Cleaner string formatting for the function call construction.


1062-1062: LGTM: Removed redundant explicit base class.

Consistent Python 3 style.

tilelang/jit/adapter/utils.py (2)

4-4: LGTM: Modernized type hints.

Updated to use PEP 604 union syntax (|) and built-in tuple instead of Union and Tuple, consistent with Python 3.10+ conventions.

Also applies to: 68-72, 110-110


75-76: LGTM: Added input validation.

Early validation ensures model_type is one of the expected values, providing clear error messages.

tilelang/language/proxy.py (5)

4-4: LGTM: Updated imports.

Removed unnecessary imports as the built-in tuple is now used instead of Tuple.


146-146: LGTM: Modernized type hint.

Using built-in tuple instead of Tuple from typing module, consistent with Python 3.9+ style.


154-154: LGTM: Modernized type hint.

Using PEP 604 union syntax (|) and built-in tuple, consistent with modern Python conventions.


175-176: LGTM: Modernized type hints.

Consistent use of built-in tuple type.


276-276: LGTM: Modernized type hint.

Using PEP 604 union syntax (str | None) instead of Optional[str].

p = get_path("README.md")
if os.path.isfile(p):
return io.open(get_path("README.md"), "r", encoding="utf-8").read()
return open(get_path("README.md"), encoding="utf-8").read() # noqa: SIM115
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Remove unused noqa directive.

The # noqa: SIM115 directive is unnecessary since SIM115 is not enabled in the ruff configuration. The static analysis tool correctly flagged this.

Apply this diff:

-    return open(get_path("README.md"), encoding="utf-8").read()  # noqa: SIM115
+    return open(get_path("README.md"), encoding="utf-8").read()
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
return open(get_path("README.md"), encoding="utf-8").read() # noqa: SIM115
return open(get_path("README.md"), encoding="utf-8").read()
🧰 Tools
🪛 Ruff (0.13.3)

312-312: Unused noqa directive (non-enabled: SIM115)

Remove unused noqa directive

(RUF100)

🤖 Prompt for AI Agents
In setup.py around line 312, the trailing "# noqa: SIM115" is unnecessary
because SIM115 is not enabled in ruff; remove the " # noqa: SIM115" suffix from
the return statement so the line simply reads return open(get_path("README.md"),
encoding="utf-8").read(), and save the file.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py (1)

1-1: Fix formatting issues before merging.

Pipeline failures indicate that this file does not match the project's formatting standards:

  • CI: "Reformatted by formatter. Changes not staged for commit."
  • CI Test on Metal: "Reformatted files. Please review and stage the changes."
  • CI Test on AMD: "clang-format reformatted file. Please review and stage the changes."

Please run the project's formatter (likely yapf based on pyproject.toml) on this file and commit the changes.

🧹 Nitpick comments (1)
examples/flash_attention/example_mha_bwd_bhsd.py (1)

264-264: Consider reverting to list comprehension for clarity.

While unpacking a generator expression is syntactically valid, it's less common and arguably less clear than the list comprehension. For 5 elements, the memory/performance benefit is negligible.

If you prefer the more conventional pattern, apply this diff:

-        do, q, k, v, o = (maybe_contiguous(x) for x in (do, q, k, v, o))
+        do, q, k, v, o = [maybe_contiguous(x) for x in (do, q, k, v, o)]
📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a13cde2 and 0973611.

📒 Files selected for processing (72)
  • docs/conf.py (1 hunks)
  • examples/amd/example_amd_flash_attn_bwd.py (2 hunks)
  • examples/attention_sink/example_gqa_sink_bwd_bhsd.py (3 hunks)
  • examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py (2 hunks)
  • examples/attention_sink/example_mha_sink_bwd_bhsd.py (3 hunks)
  • examples/attention_sink/example_mha_sink_fwd_bhsd.py (2 hunks)
  • examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py (2 hunks)
  • examples/bitnet-1.58b/configuration_bitnet.py (0 hunks)
  • examples/bitnet-1.58b/eval_ppl.py (1 hunks)
  • examples/bitnet-1.58b/maint/create_bitblas_ckpt.py (1 hunks)
  • examples/bitnet-1.58b/modeling_bitnet.py (1 hunks)
  • examples/bitnet-1.58b/tokenization_bitnet.py (0 hunks)
  • examples/bitnet-1.58b/utils_quant.py (1 hunks)
  • examples/bitnet-1.58b/vllm_workspace/conftest.py (1 hunks)
  • examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py (2 hunks)
  • examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py (2 hunks)
  • examples/cast/example_group_per_split_token_cast_to_fp8.py (1 hunks)
  • examples/cast/example_per_token_cast_to_fp8.py (2 hunks)
  • examples/deepseek_mla/example_mla_decode_paged.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_bf16_fp4_hopper.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper_tma.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_fp4_hopper.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_groupedgemm_bf16_mxfp4_hopper.py (1 hunks)
  • examples/flash_attention/example_gqa_bwd.py (2 hunks)
  • examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (2 hunks)
  • examples/flash_attention/example_gqa_fwd_bshd.py (1 hunks)
  • examples/flash_attention/example_gqa_fwd_bshd_wgmma_pipelined.py (1 hunks)
  • examples/flash_attention/example_mha_bwd.py (2 hunks)
  • examples/flash_attention/example_mha_bwd_bhsd.py (2 hunks)
  • examples/flash_attention/example_mha_bwd_wgmma_pipelined.py (2 hunks)
  • examples/flash_attention/example_mha_fwd_bhsd.py (1 hunks)
  • examples/flash_attention/example_mha_fwd_bhsd_wgmma_pipelined.py (1 hunks)
  • examples/flash_attention/example_mha_fwd_bshd.py (1 hunks)
  • examples/flash_attention/example_mha_fwd_bshd_wgmma_pipelined.py (1 hunks)
  • examples/flash_decoding/example_gqa_decode.py (1 hunks)
  • examples/flash_decoding/example_mha_inference.py (1 hunks)
  • examples/grouped_gemm/example_grouped_gemm_bwd.py (1 hunks)
  • examples/hadamard_transform/example_hadamard.py (1 hunks)
  • examples/linear_attention/example_mamba_chunk_scan.py (1 hunks)
  • examples/linear_attention/example_mamba_chunk_state.py (2 hunks)
  • examples/minference/example_vertical_slash_sparse_attn.py (1 hunks)
  • examples/norm/rms_norm.py (1 hunks)
  • pyproject.toml (1 hunks)
  • setup.py (6 hunks)
  • testing/python/kernel/test_tilelang_kernel_gemm.py (1 hunks)
  • testing/python/kernel/test_tilelang_kernel_gemm_simt.py (1 hunks)
  • testing/python/language/test_tilelang_language_pipeline.py (1 hunks)
  • tilelang/autotuner/param.py (3 hunks)
  • tilelang/cache/kernel_cache.py (1 hunks)
  • tilelang/carver/arch/cuda.py (1 hunks)
  • tilelang/carver/arch/metal.py (1 hunks)
  • tilelang/carver/roller/bestfit.py (1 hunks)
  • tilelang/carver/roller/hint.py (1 hunks)
  • tilelang/carver/roller/node.py (4 hunks)
  • tilelang/carver/roller/rasterization.py (1 hunks)
  • tilelang/carver/roller/shape_inference/common.py (2 hunks)
  • tilelang/carver/roller/shape_inference/tir.py (3 hunks)
  • tilelang/contrib/hipcc.py (1 hunks)
  • tilelang/intrinsics/mfma_macro_generator.py (1 hunks)
  • tilelang/intrinsics/mma_macro_generator.py (2 hunks)
  • tilelang/intrinsics/wgmma_macro_generator.py (1 hunks)
  • tilelang/jit/adapter/cython/adapter.py (3 hunks)
  • tilelang/jit/adapter/libgen.py (1 hunks)
  • tilelang/jit/adapter/utils.py (3 hunks)
  • tilelang/jit/adapter/wrapper.py (6 hunks)
  • tilelang/jit/kernel.py (1 hunks)
  • tilelang/language/proxy.py (4 hunks)
  • tilelang/quantize/lop3.py (1 hunks)
  • tilelang/quantize/quantization.py (2 hunks)
  • tilelang/tileop/gemm/gemm_base.py (1 hunks)
  • tilelang/version.py (1 hunks)
💤 Files with no reviewable changes (2)
  • examples/bitnet-1.58b/tokenization_bitnet.py
  • examples/bitnet-1.58b/configuration_bitnet.py
🧰 Additional context used
🧬 Code graph analysis (23)
examples/flash_attention/example_mha_fwd_bhsd.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
tilelang/intrinsics/mma_macro_generator.py (2)
tilelang/intrinsics/wgmma_macro_generator.py (1)
  • TensorCoreIntrinEmitter (64-519)
tilelang/utils/language.py (1)
  • is_fragment (68-78)
examples/flash_decoding/example_gqa_decode.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
tilelang/intrinsics/wgmma_macro_generator.py (1)
tilelang/utils/language.py (1)
  • is_fragment (68-78)
examples/flash_decoding/example_mha_inference.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_mha_fwd_bshd.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_mha_fwd_bhsd_wgmma_pipelined.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/dequantize_gemm/example_dequant_gemm_fp4_hopper.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_gqa_bwd.py (2)
examples/amd/example_amd_flash_attn_bwd.py (2)
  • maybe_contiguous (242-245)
  • run1 (340-341)
examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (2)
  • maybe_contiguous (399-402)
  • run1 (536-537)
examples/flash_attention/example_gqa_fwd_bshd_wgmma_pipelined.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/attention_sink/example_mha_sink_bwd_bhsd.py (3)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (1)
  • maybe_contiguous (356-359)
examples/flash_attention/example_mha_bwd.py (1)
  • maybe_contiguous (258-261)
examples/minference/example_vertical_slash_sparse_attn.py (1)
tilelang/language/builtin.py (1)
  • mbarrier_wait_parity (172-219)
examples/flash_attention/example_mha_fwd_bshd_wgmma_pipelined.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/linear_attention/example_mamba_chunk_scan.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
tilelang/jit/adapter/utils.py (1)
tilelang/language/ast/ir.py (1)
  • target (1682-1713)
examples/flash_attention/example_mha_bwd_wgmma_pipelined.py (1)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/amd/example_amd_flash_attn_bwd.py (1)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (1)
  • maybe_contiguous (356-359)
examples/flash_attention/example_mha_bwd_bhsd.py (3)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (1)
  • maybe_contiguous (356-359)
examples/flash_attention/example_mha_bwd.py (1)
  • maybe_contiguous (258-261)
examples/linear_attention/example_mamba_chunk_state.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_gqa_fwd_bshd.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (2)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/attention_sink/example_mha_sink_bwd_bhsd.py (2)
  • maybe_contiguous (364-367)
  • tl_bwd (488-489)
examples/flash_attention/example_mha_bwd.py (2)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (1)
examples/flash_attention/example_gqa_bwd.py (2)
  • maybe_contiguous (377-380)
  • run1 (514-515)
🪛 GitHub Actions: CI
examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py

[error] 1-1: Reformatted by formatter. Changes not staged for commit.

docs/conf.py

[error] 1-1: Reformatted by formatter. Changes not staged for commit.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py

[error] 1-1: Reformatted by formatter. Changes not staged for commit.

🪛 GitHub Actions: CI Test on AMD
examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py

[error] 1-1: clang-format reformatted file. Please review and stage the changes.

docs/conf.py

[error] 1-1: clang-format reformatted file. Please review and stage the changes.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py

[error] 1-1: clang-format reformatted file. Please review and stage the changes.

🪛 GitHub Actions: CI Test on Metal
examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py

[error] 1-1: Reformatted files. Please review and stage the changes. Changes not staged for commit.

docs/conf.py

[error] 1-1: Reformatted files. Please review and stage the changes. Changes not staged for commit.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py

[error] 1-1: Reformatted files. Please review and stage the changes. Changes not staged for commit.

🪛 Ruff (0.13.3)
tilelang/quantize/lop3.py

1189-1189: Avoid specifying long messages outside the exception class

(TRY003)

tilelang/carver/roller/shape_inference/tir.py

354-354: Prefer TypeError exception for invalid type

(TRY004)


354-354: Create your own exception

(TRY002)


354-354: Avoid specifying long messages outside the exception class

(TRY003)

tilelang/carver/roller/node.py

304-304: Use of functools.lru_cache or functools.cache on methods can lead to memory leaks

(B019)


421-421: Use of functools.lru_cache or functools.cache on methods can lead to memory leaks

(B019)

setup.py

312-312: Unused noqa directive (non-enabled: SIM115)

Remove unused noqa directive

(RUF100)

🔇 Additional comments (70)
examples/flash_decoding/example_mha_inference.py (1)

321-325: F-string conversion looks good

Formatting and precision stay the same, readability improves.

testing/python/kernel/test_tilelang_kernel_gemm.py (1)

98-99: Redundant parentheses cleanup looks good.

The simplified expression preserves behavior and keeps the intent clear.

examples/cast/example_group_per_split_token_cast_to_fp8.py (1)

202-205: LGTM on the f-string swap

Straightforward modernization—f-strings keep the same output while matching the new Ruff style expectations.

examples/hadamard_transform/example_hadamard.py (1)

154-154: LGTM! Modern f-string formatting applied correctly.

The conversion from .format() to f-string is correct and improves readability while maintaining identical output.

examples/cast/example_per_token_cast_to_fp8.py (1)

103-103: LGTM! Clean modernization to f-strings.

The migration from .format() to f-strings is correct and aligns with the PR objective of enabling pyupgrade linter. The formatting specifications are preserved accurately.

Also applies to: 105-105, 117-117

examples/norm/rms_norm.py (1)

74-76: LGTM! Clean modernization to f-strings.

The conversion from format strings to f-strings is correct and improves readability. The formatting specifiers are preserved accurately.

examples/bitnet-1.58b/eval_ppl.py (1)

57-57: LGTM! F-string modernization improves readability.

The conversion to an f-string is correct and aligns with modern Python best practices. It also maintains consistency with the existing f-string usage at line 53.

testing/python/language/test_tilelang_language_pipeline.py (1)

106-107: LGTM! Clean formatting improvement.

Removing the redundant outer parentheses improves readability without affecting functionality. The tfloat32 conversion logic remains correct.

tilelang/carver/roller/rasterization.py (1)

92-92: LGTM! Clean f-string conversion.

The modernization from .format() to an f-string is correct and aligns with the PR's objective to enable pyupgrade linting.

examples/dequantize_gemm/example_dequant_groupedgemm_bf16_mxfp4_hopper.py (1)

493-494: LGTM! Clean modernization to f-strings.

The conversion from .format() to f-strings is correct and improves readability. The formatting specifiers are preserved accurately, and the changes align with the PR's objective of modernizing Python style across the codebase.

testing/python/kernel/test_tilelang_kernel_gemm_simt.py (1)

109-109: LGTM! Redundant parentheses removed.

The removal of redundant parentheses around block_K // micro_size_k improves code clarity and aligns with the usage of T.serial elsewhere in the file (lines 99, 110, 115, 121, 126). This change is part of the pyupgrade linter enforcement and has no functional impact.

examples/flash_attention/example_mha_fwd_bshd_wgmma_pipelined.py (1)

198-202: LGTM! Clean f-string conversion.

The print statements have been correctly modernized to use f-strings, improving readability while maintaining identical output formatting. The conversions follow Python best practices and align with the PR's objective to enable pyupgrade linting.

examples/bitnet-1.58b/utils_quant.py (1)

219-219: LGTM! Modern Python 3 super() call.

The modernization from super(BitLinear, self).__init__(...) to super().__init__(...) is correct and follows Python 3 best practices. This change is consistent with the existing code in BitLinearBitBLAS.__init__ (line 44) and aligns with the PR's objective to modernize Python style.

examples/linear_attention/example_mamba_chunk_state.py (2)

43-43: LGTM! Redundant parentheses removed.

The extra parentheses around the subtraction expression are unnecessary and have been correctly removed without affecting functionality.


171-175: LGTM! Print statements modernized to f-strings.

The conversion from .format() to f-strings improves readability and aligns with modern Python style. The formatting specifiers and output values remain identical.

examples/flash_attention/example_gqa_fwd_bshd_wgmma_pipelined.py (2)

225-226: LGTM! F-string conversion is correct.

The conversion from .format() to f-strings is properly done and aligns with the PR objective of modernizing Python style.


228-229: LGTM! Consistent f-string modernization.

The f-string conversion matches the pattern used for the reference benchmark output above, maintaining consistency throughout the file.

examples/flash_attention/example_mha_bwd_bhsd.py (1)

342-346: LGTM! Good modernization to f-strings.

The conversion from .format() to f-strings improves readability and aligns with modern Python style.

examples/bitnet-1.58b/vllm_workspace/conftest.py (1)

37-37: LGTM! Clean modernization.

Removing the explicit "r" mode is appropriate since it's Python's default for open(). This change aligns with the pyupgrade linter recommendations.

tilelang/carver/arch/metal.py (1)

1-1: LGTM! Enables modern type hint syntax.

Adding from __future__ import annotations is appropriate for this module. It enables the modern union syntax (Target | str on line 12) and improves forward reference handling during type checking.

examples/bitnet-1.58b/modeling_bitnet.py (1)

1515-1515: LGTM! Consistent with file I/O modernization.

Removing the explicit "r" mode follows the same pattern as other file I/O updates in this PR and aligns with pyupgrade recommendations.

tilelang/intrinsics/wgmma_macro_generator.py (1)

408-408: LGTM! F-string improves readability.

The conversion from .format() to an f-string is correct and aligns with modern Python style guidelines enforced by pyupgrade.

examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper_tma.py (1)

553-554: LGTM! Print statements modernized.

The f-string conversions for the latency and TFlops output are correct and improve code readability while maintaining identical output formatting.

tilelang/carver/roller/bestfit.py (1)

20-20: LGTM! Cleaner repr implementation.

The f-string conversion in the __repr__ method is correct and makes the code more readable while preserving the exact string representation.

examples/flash_attention/example_mha_fwd_bhsd_wgmma_pipelined.py (1)

211-215: LGTM! Benchmark output formatting improved.

All four print statements have been correctly converted to f-strings, maintaining the same output format while improving code clarity and consistency with the rest of the PR.

examples/bitnet-1.58b/maint/create_bitblas_ckpt.py (1)

88-88: LGTM! Consistent file I/O modernization.

Removing the explicit "r" mode is appropriate and aligns with the broader PR pattern of modernizing file I/O operations across the codebase.

pyproject.toml (1)

34-34: LGTM! Pyupgrade rules enabled.

Enabling the UP (pyupgrade) and FA102 rules aligns with the PR objective to prevent issues by modernizing Python syntax across the codebase.

examples/flash_attention/example_mha_fwd_bhsd.py (1)

206-210: LGTM! F-string conversion improves readability.

The conversion from .format() to f-strings is a standard Python modernization that improves code readability while maintaining identical functionality.

tilelang/carver/roller/hint.py (1)

150-150: LGTM! Python 3 class declaration modernization.

Removing explicit object inheritance is correct for Python 3, where all classes implicitly inherit from object. This is a standard pyupgrade modernization with no behavioral change.

tilelang/autotuner/param.py (3)

249-249: LGTM! Default file mode simplification.

Omitting the explicit "r" mode is correct since text-read mode is the default for open(). This is a standard pyupgrade modernization.


318-318: LGTM! Default file mode simplification.

Omitting the explicit "r" mode is correct since text-read mode is the default for open().


330-330: LGTM! Default file mode simplification.

Omitting the explicit "r" mode is correct since text-read mode is the default for open().

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py (2)

183-183: LGTM! Modern super() call.

The modern super() syntax without explicit class and self arguments is the recommended Python 3 pattern and is functionally equivalent.


357-362: LGTM! F-string conversions improve readability.

The conversion to f-strings modernizes the code while maintaining identical output.

tilelang/jit/kernel.py (1)

20-20: LGTM! Python 3 class declaration modernization.

Removing explicit object inheritance is correct for Python 3. This is a standard pyupgrade modernization with no behavioral change.

tilelang/jit/adapter/libgen.py (1)

32-32: LGTM! Python 3 class declaration modernization.

Removing explicit object inheritance is correct for Python 3. This is a standard pyupgrade modernization with no behavioral change.

examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (2)

404-404: LGTM! Generator expression for unpacking is more efficient.

Using a generator expression instead of a list comprehension for tuple unpacking is more memory-efficient. Python will consume the generator during unpacking to assign all values correctly.


542-546: LGTM! F-string conversions improve readability.

The conversion from .format() to f-strings is a standard Python modernization that improves code readability.

tilelang/language/proxy.py (1)

4-4: LGTM! Type hints modernized to Python 3.10+ syntax.

The changes correctly update type hints to use:

  • Built-in tuple instead of typing.Tuple
  • Union syntax X | Y instead of Union[X, Y]
  • Union syntax X | None instead of Optional[X]

These align with PEP 604 and PEP 585 best practices.

Also applies to: 146-146, 154-154, 175-176, 276-276

tilelang/jit/adapter/utils.py (1)

4-4: LGTM! Type hints modernized and validation improved.

The changes correctly:

  • Replace Dict with built-in dict[...]
  • Use union syntax X | Y instead of Union[X, Y]
  • Use X | None instead of Optional[X]

The early validation for model_type at lines 74-76 is a good addition that provides clearer error messages.

Also applies to: 67-72, 110-110

examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py (2)

3-3: LGTM! Future annotations import added.

This enables postponed evaluation of annotations (PEP 563), which is required for the modern type hints used elsewhere in the codebase.


423-427: LGTM! Print statements modernized to f-strings.

The f-string conversions preserve the exact formatting (.2f) and improve readability. No functional changes.

examples/flash_decoding/example_gqa_decode.py (1)

476-480: LGTM! Print statements modernized to f-strings.

The f-string conversions preserve the exact formatting (.2f) and improve code readability.

tilelang/version.py (1)

24-24: LGTM! Explicit read mode removed.

The explicit "r" mode is redundant as it's the default for open(). This modernization aligns with Python best practices.

docs/conf.py (1)

5-5: LGTM! String formatting and file opening modernized.

The changes correctly:

  • Convert % formatting to f-string for better readability
  • Remove redundant explicit "r" mode (default for open())

However, note the pipeline failures indicating formatting issues that need to be staged.

Please address the pipeline failures mentioned in the CI:

  • "Reformatted by formatter. Changes not staged for commit."
  • "clang-format reformatted file. Please review and stage the changes."

These appear to be formatting issues that need to be staged for commit.

Also applies to: 8-8

examples/attention_sink/example_gqa_sink_bwd_bhsd.py (3)

2-2: LGTM! Future annotations import added.

This enables postponed evaluation of annotations (PEP 563), supporting the modern type hints used throughout the codebase.


361-361: LGTM! List comprehension changed to generator expression.

The generator expression is more memory-efficient and is safe here because the result is immediately unpacked into individual variables. This is a good optimization for transforming multiple items.


488-492: LGTM! Print statements modernized to f-strings.

The f-string conversions preserve exact formatting (.2f) and improve readability.

examples/attention_sink/example_mha_sink_fwd_bhsd.py (2)

2-2: LGTM! Future annotations import added.

This enables postponed evaluation of annotations (PEP 563), which is necessary for modern type hint syntax.


289-293: LGTM! Print statements modernized to f-strings.

The f-string conversions preserve the exact formatting (.2f) and are more Pythonic.

tilelang/cache/kernel_cache.py (1)

366-366: LGTM! Safe modernization.

Removing the explicit "r" mode is a safe Python 3 idiom since "r" is the default for open().

examples/deepseek_mla/example_mla_decode_paged.py (1)

403-404: LGTM! Modern f-string syntax.

The conversion from .format() to f-strings is a safe, idiomatic modernization with identical output.

examples/flash_attention/example_mha_bwd.py (2)

263-263: LGTM! Memory-efficient generator unpacking.

Replacing the list comprehension with a generator expression is safe here since the unpacking consumes exactly five elements. This reduces memory overhead by avoiding an intermediate list.


339-343: LGTM! Modern f-string syntax.

The conversion from .format() to f-strings is a safe, idiomatic modernization with identical output.

tilelang/tileop/gemm/gemm_base.py (1)

12-12: LGTM! Implicit object inheritance.

In Python 3, all classes implicitly inherit from object, so removing the explicit base is a safe modernization. The @dataclass decorator and class behavior remain unchanged.

examples/amd/example_amd_flash_attn_bwd.py (2)

247-247: LGTM! Memory-efficient generator unpacking.

Replacing the list comprehension with a generator expression is safe here since the unpacking consumes exactly five elements. This reduces memory overhead by avoiding an intermediate list.


346-350: LGTM! Modern f-string syntax.

The conversion from .format() to f-strings is a safe, idiomatic modernization with identical output.

setup.py (3)

140-140: LGTM! Safe modernization.

Removing the explicit "r" mode is a safe Python 3 idiom since "r" is the default for open().


171-171: LGTM! Safe modernization.

Removing the explicit "r" mode is a safe Python 3 idiom since "r" is the default for open().


711-740: LGTM! Safe modernization.

Removing the explicit "r" mode from text file opens is a safe Python 3 idiom since "r" is the default. The behavior remains identical.

examples/attention_sink/example_mha_sink_bwd_bhsd.py (3)

2-2: LGTM! Forward-compatible annotations.

Adding from __future__ import annotations enables PEP 563 (postponed evaluation of annotations), which is forward-compatible and allows using newer annotation syntax. This is a safe, future-proof change.


369-369: LGTM! Memory-efficient generator unpacking.

Replacing the list comprehension with a generator expression is safe here since the unpacking consumes exactly six elements. This reduces memory overhead by avoiding an intermediate list.


492-496: LGTM! Modern f-string syntax.

The conversion from .format() to f-strings is a safe, idiomatic modernization with identical output.

examples/flash_attention/example_gqa_bwd.py (2)

382-382: LGTM! Memory-efficient generator unpacking.

Replacing the list comprehension with a generator expression is safe here since the unpacking consumes exactly five elements. This reduces memory overhead by avoiding an intermediate list.


520-524: LGTM! Modern f-string syntax.

The conversion from .format() to f-strings is a safe, idiomatic modernization with identical output.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py (3)

196-196: LGTM! Modern super() syntax.

The change from super(SparseFlashAttn, self).__init__() to super().__init__() follows Python 3 best practices and is the idiomatic way to call parent class constructors.


369-369: LGTM! f-string conversion.

The change from .format() to f-string improves readability and aligns with modern Python style.


372-372: LGTM! f-string conversion with multiple expressions.

The conversion correctly maps all format arguments to f-string expressions, improving readability.

tilelang/jit/adapter/cython/adapter.py (1)

105-105: LGTM! Default file mode usage.

Removing the explicit "r" mode is correct since Python defaults to text read mode. This modernization aligns with pyupgrade linter rules and makes the code cleaner without changing behavior.

Also applies to: 118-118, 135-135

examples/grouped_gemm/example_grouped_gemm_bwd.py (1)

139-139: Confirm generator-expression change
The switch from [maybe_contiguous(x) for x in …] to (maybe_contiguous(x) for x in …) for immediate unpacking isn’t a known pyupgrade rule and offers no benefit. Verify this refactor is intentional and actually produced by your pyupgrade/ruff setup (e.g., share your lint output or rule config).

p = get_path("README.md")
if os.path.isfile(p):
return io.open(get_path("README.md"), "r", encoding="utf-8").read()
return open(get_path("README.md"), encoding="utf-8").read() # noqa: SIM115
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Remove unused noqa directive.

The # noqa: SIM115 directive is flagged as unused because SIM115 (use context manager) is not triggered by this code. The directive should be removed.

Apply this diff:

-        return open(get_path("README.md"), encoding="utf-8").read()  # noqa: SIM115
+        return open(get_path("README.md"), encoding="utf-8").read()
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
return open(get_path("README.md"), encoding="utf-8").read() # noqa: SIM115
return open(get_path("README.md"), encoding="utf-8").read()
🧰 Tools
🪛 Ruff (0.13.3)

312-312: Unused noqa directive (non-enabled: SIM115)

Remove unused noqa directive

(RUF100)

🤖 Prompt for AI Agents
In setup.py around line 312, the return statement uses an unnecessary noqa
directive; remove the trailing "  # noqa: SIM115" from the line so the file
reads simply return open(get_path("README.md"), encoding="utf-8").read(),
eliminating the unused lint suppression.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (5)
examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py (1)

357-357: Good f-string conversion with minor optimization opportunity.

The f-string conversion improves readability. Consider simplifying further by eliminating the concatenation:

-    print(name + f"  all_close={all_close}")
+    print(f"{name}  all_close={all_close}")
examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py (1)

369-369: Good f-string conversion with minor optimization opportunity.

The f-string conversion improves readability. Consider simplifying further by eliminating the concatenation:

-    print(name + f"  all_close={all_close}")
+    print(f"{name}  all_close={all_close}")
examples/grouped_gemm/example_grouped_gemm_bwd.py (1)

139-139: Consider using map() for better idiomaticity.

The generator expression works correctly but is unconventional for immediate unpacking. The map() builtin is the standard Python pattern for applying a function to multiple values.

Apply this diff to use the more idiomatic map():

-        A, B, batch_sizes = (maybe_contiguous(x) for x in (A, B, batch_sizes))
+        A, B, batch_sizes = map(maybe_contiguous, (A, B, batch_sizes))
setup.py (1)

312-312: Remove the unused noqa directive.

The # noqa: SIM115 comment is no longer needed since the code has been updated to address the linting issue.

Apply this diff:

-        return open(get_path("README.md"), encoding="utf-8").read()  # noqa: SIM115
+        return open(get_path("README.md"), encoding="utf-8").read()
tilelang/carver/roller/shape_inference/tir.py (1)

354-354: F-string migration looks good; consider using TypeError for type errors.

The f-string conversion is correct. As an optional improvement, consider raising TypeError instead of generic Exception when encountering unexpected types, as this provides clearer intent.

Apply this diff if desired:

-        raise Exception(f"Unhandled node type in walk_indice(): {expr}")
+        raise TypeError(f"Unhandled node type in walk_indice(): {type(expr).__name__}")
📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a13cde2 and bfe2ff8.

📒 Files selected for processing (72)
  • docs/conf.py (1 hunks)
  • examples/amd/example_amd_flash_attn_bwd.py (2 hunks)
  • examples/attention_sink/example_gqa_sink_bwd_bhsd.py (3 hunks)
  • examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py (2 hunks)
  • examples/attention_sink/example_mha_sink_bwd_bhsd.py (3 hunks)
  • examples/attention_sink/example_mha_sink_fwd_bhsd.py (2 hunks)
  • examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py (2 hunks)
  • examples/bitnet-1.58b/configuration_bitnet.py (0 hunks)
  • examples/bitnet-1.58b/eval_ppl.py (1 hunks)
  • examples/bitnet-1.58b/maint/create_bitblas_ckpt.py (1 hunks)
  • examples/bitnet-1.58b/modeling_bitnet.py (1 hunks)
  • examples/bitnet-1.58b/tokenization_bitnet.py (0 hunks)
  • examples/bitnet-1.58b/utils_quant.py (1 hunks)
  • examples/bitnet-1.58b/vllm_workspace/conftest.py (1 hunks)
  • examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py (2 hunks)
  • examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py (2 hunks)
  • examples/cast/example_group_per_split_token_cast_to_fp8.py (1 hunks)
  • examples/cast/example_per_token_cast_to_fp8.py (2 hunks)
  • examples/deepseek_mla/example_mla_decode_paged.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_bf16_fp4_hopper.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper_tma.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_gemm_fp4_hopper.py (1 hunks)
  • examples/dequantize_gemm/example_dequant_groupedgemm_bf16_mxfp4_hopper.py (1 hunks)
  • examples/flash_attention/example_gqa_bwd.py (2 hunks)
  • examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (2 hunks)
  • examples/flash_attention/example_gqa_fwd_bshd.py (1 hunks)
  • examples/flash_attention/example_gqa_fwd_bshd_wgmma_pipelined.py (1 hunks)
  • examples/flash_attention/example_mha_bwd.py (2 hunks)
  • examples/flash_attention/example_mha_bwd_bhsd.py (2 hunks)
  • examples/flash_attention/example_mha_bwd_wgmma_pipelined.py (2 hunks)
  • examples/flash_attention/example_mha_fwd_bhsd.py (1 hunks)
  • examples/flash_attention/example_mha_fwd_bhsd_wgmma_pipelined.py (1 hunks)
  • examples/flash_attention/example_mha_fwd_bshd.py (1 hunks)
  • examples/flash_attention/example_mha_fwd_bshd_wgmma_pipelined.py (1 hunks)
  • examples/flash_decoding/example_gqa_decode.py (1 hunks)
  • examples/flash_decoding/example_mha_inference.py (1 hunks)
  • examples/grouped_gemm/example_grouped_gemm_bwd.py (1 hunks)
  • examples/hadamard_transform/example_hadamard.py (1 hunks)
  • examples/linear_attention/example_mamba_chunk_scan.py (1 hunks)
  • examples/linear_attention/example_mamba_chunk_state.py (2 hunks)
  • examples/minference/example_vertical_slash_sparse_attn.py (1 hunks)
  • examples/norm/rms_norm.py (1 hunks)
  • pyproject.toml (1 hunks)
  • setup.py (6 hunks)
  • testing/python/kernel/test_tilelang_kernel_gemm.py (1 hunks)
  • testing/python/kernel/test_tilelang_kernel_gemm_simt.py (1 hunks)
  • testing/python/language/test_tilelang_language_pipeline.py (1 hunks)
  • tilelang/autotuner/param.py (3 hunks)
  • tilelang/cache/kernel_cache.py (1 hunks)
  • tilelang/carver/arch/cuda.py (1 hunks)
  • tilelang/carver/arch/metal.py (1 hunks)
  • tilelang/carver/roller/bestfit.py (1 hunks)
  • tilelang/carver/roller/hint.py (1 hunks)
  • tilelang/carver/roller/node.py (4 hunks)
  • tilelang/carver/roller/rasterization.py (1 hunks)
  • tilelang/carver/roller/shape_inference/common.py (2 hunks)
  • tilelang/carver/roller/shape_inference/tir.py (3 hunks)
  • tilelang/contrib/hipcc.py (1 hunks)
  • tilelang/intrinsics/mfma_macro_generator.py (1 hunks)
  • tilelang/intrinsics/mma_macro_generator.py (2 hunks)
  • tilelang/intrinsics/wgmma_macro_generator.py (1 hunks)
  • tilelang/jit/adapter/cython/adapter.py (3 hunks)
  • tilelang/jit/adapter/libgen.py (1 hunks)
  • tilelang/jit/adapter/utils.py (3 hunks)
  • tilelang/jit/adapter/wrapper.py (6 hunks)
  • tilelang/jit/kernel.py (1 hunks)
  • tilelang/language/proxy.py (4 hunks)
  • tilelang/quantize/lop3.py (1 hunks)
  • tilelang/quantize/quantization.py (2 hunks)
  • tilelang/tileop/gemm/gemm_base.py (1 hunks)
  • tilelang/version.py (1 hunks)
💤 Files with no reviewable changes (2)
  • examples/bitnet-1.58b/configuration_bitnet.py
  • examples/bitnet-1.58b/tokenization_bitnet.py
🧰 Additional context used
🧬 Code graph analysis (25)
examples/flash_attention/example_mha_fwd_bhsd_wgmma_pipelined.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/minference/example_vertical_slash_sparse_attn.py (1)
tilelang/language/builtin.py (1)
  • mbarrier_wait_parity (172-219)
examples/dequantize_gemm/example_dequant_gemm_fp4_hopper.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/linear_attention/example_mamba_chunk_scan.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
tilelang/jit/adapter/utils.py (1)
tilelang/language/ast/ir.py (1)
  • target (1682-1713)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (2)
examples/attention_sink/example_mha_sink_bwd_bhsd.py (1)
  • maybe_contiguous (364-367)
examples/flash_attention/example_gqa_bwd.py (1)
  • maybe_contiguous (377-380)
examples/flash_decoding/example_gqa_decode.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_mha_bwd.py (2)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
tilelang/carver/roller/shape_inference/common.py (1)
tilelang/carver/roller/shape_inference/tir.py (2)
  • Statement (7-43)
  • InputShapeInference (169-318)
examples/linear_attention/example_mamba_chunk_state.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
tilelang/intrinsics/mma_macro_generator.py (1)
tilelang/utils/language.py (1)
  • is_fragment (68-78)
examples/flash_attention/example_mha_fwd_bshd.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/grouped_gemm/example_grouped_gemm_bwd.py (1)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/flash_decoding/example_mha_inference.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_gqa_fwd_bshd.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_mha_fwd_bshd_wgmma_pipelined.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_gqa_fwd_bshd_wgmma_pipelined.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_gqa_bwd.py (1)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/flash_attention/example_mha_fwd_bhsd.py (1)
tilelang/profiler/__init__.py (1)
  • do_bench (218-281)
examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (1)
examples/flash_attention/example_gqa_bwd.py (2)
  • maybe_contiguous (377-380)
  • run1 (514-515)
examples/amd/example_amd_flash_attn_bwd.py (1)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (1)
  • maybe_contiguous (356-359)
examples/flash_attention/example_mha_bwd_wgmma_pipelined.py (1)
examples/amd/example_amd_flash_attn_bwd.py (1)
  • maybe_contiguous (242-245)
examples/attention_sink/example_mha_sink_bwd_bhsd.py (1)
examples/flash_attention/example_mha_bwd.py (1)
  • maybe_contiguous (258-261)
examples/flash_attention/example_mha_bwd_bhsd.py (1)
examples/amd/example_amd_flash_attn_bwd.py (2)
  • maybe_contiguous (242-245)
  • run1 (340-341)
tilelang/intrinsics/wgmma_macro_generator.py (1)
tilelang/utils/language.py (1)
  • is_fragment (68-78)
🪛 Ruff (0.13.3)
tilelang/carver/roller/shape_inference/tir.py

354-354: Prefer TypeError exception for invalid type

(TRY004)


354-354: Create your own exception

(TRY002)


354-354: Avoid specifying long messages outside the exception class

(TRY003)

setup.py

312-312: Unused noqa directive (non-enabled: SIM115)

Remove unused noqa directive

(RUF100)

tilelang/quantize/lop3.py

1189-1189: Avoid specifying long messages outside the exception class

(TRY003)

tilelang/carver/roller/node.py

304-304: Use of functools.lru_cache or functools.cache on methods can lead to memory leaks

(B019)


421-421: Use of functools.lru_cache or functools.cache on methods can lead to memory leaks

(B019)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
  • GitHub Check: build-test-amd
  • GitHub Check: format-check
🔇 Additional comments (98)
testing/python/language/test_tilelang_language_pipeline.py (1)

106-107: No-op cleanup looks good.

Dropping the redundant parentheses keeps the expression readable with identical semantics while satisfying the linter. ✅

examples/minference/example_vertical_slash_sparse_attn.py (1)

209-209: Redundant parentheses removal looks good.

The simplified parity expression keeps identical behavior while matching the linter’s style expectations.

examples/bitnet-1.58b/eval_ppl.py (1)

57-57: LGTM! Clean modernization to f-string.

The conversion from .format() to f-string is correct and improves readability. This change also enhances consistency, as line 53 already uses f-strings for similar output formatting.

tilelang/contrib/hipcc.py (1)

57-57: LGTM! Clean string formatting modernization.

The conversion from % formatting to f-string is correct and improves readability. Since target_format is validated to be "hsaco" on line 54-55, both the old and new formatting produce identical output.

examples/cast/example_per_token_cast_to_fp8.py (1)

103-103: LGTM! Clean modernization to f-strings.

The conversion from .format() to f-strings is correct and aligns with Python 3.6+ best practices. F-strings are more readable and slightly more performant than the older string formatting method.

Also applies to: 105-105, 117-117

examples/flash_decoding/example_gqa_decode.py (1)

476-480: F-string conversions are equivalent

Switchover keeps the formatting semantics intact while aligning with the pyupgrade style target.

tilelang/quantize/quantization.py (2)

219-227: No behavior change in prefix assembly.

Parentheses cleanup keeps the same precedence, so the packed bf16 exponent path remains identical. Looks good.


231-236: Bitwise OR chaining still correct.

Associativity of | means the expanded form matches the prior grouped expression; the e4m3 conversion stays stable.

examples/cast/example_group_per_split_token_cast_to_fp8.py (2)

202-202: LGTM! Clean f-string migration.

The conversion from .format() to f-string is correct and improves readability, aligning with the PR objective to enable pyupgrade linter rules.


205-205: LGTM! Clean f-string migration.

The conversion from .format() to f-string is correct and improves readability, consistent with the modernization effort.

examples/dequantize_gemm/example_dequant_groupedgemm_bf16_mxfp4_hopper.py (1)

493-494: LGTM! Clean modernization to f-strings.

The conversion from .format() to f-strings is correct and preserves the formatting specification (.2f). This change improves code consistency with the existing f-string usage elsewhere in the file (e.g., lines 436-437) and aligns with modern Python best practices.

examples/flash_decoding/example_mha_inference.py (1)

321-325: LGTM! String formatting modernization improves readability.

The conversion from .format() to f-strings is correct and aligns with modern Python best practices (PEP 498) and the PR's objective to enable pyupgrade linting. The formatting specifiers and calculations remain identical, ensuring no behavioral changes.

examples/dequantize_gemm/example_dequant_gemm_bf16_fp4_hopper.py (1)

436-437: LGTM: Clean f-string migration.

The conversion from .format() to f-strings improves readability and aligns with modern Python best practices.

examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper_tma.py (1)

553-554: LGTM: Consistent f-string migration.

The formatting update is correct and consistent with the broader modernization effort across the codebase.

examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper.py (1)

537-538: LGTM: F-string conversion applied correctly.

The formatting change maintains the same output while improving code quality.

examples/dequantize_gemm/example_dequant_gemm_fp4_hopper.py (1)

280-284: LGTM: F-string updates for all benchmark outputs.

All four print statements have been correctly migrated to f-strings, maintaining consistent formatting across reference and kernel benchmarks.

examples/deepseek_mla/example_mla_decode_paged.py (1)

403-404: LGTM! Clean f-string conversion.

The conversion from .format() to f-strings is correct and maintains identical output. This modernization improves readability and aligns with Python 3.6+ best practices.

examples/bitnet-1.58b/utils_quant.py (1)

219-219: LGTM! Modern super() syntax applied correctly.

The change from super(BitLinear, self).__init__(*kargs, **kwargs) to super().__init__(*kargs, **kwargs) is correct and aligns with Python 3+ best practices. The behavior remains equivalent.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_mask.py (2)

183-183: LGTM! Modern super() syntax.

The simplified super().__init__() call is the recommended Python 3 idiom, reducing boilerplate and potential errors.


362-364: LGTM! Clean multi-line f-string formatting.

The conversion from .format() to f-strings improves readability and performance while maintaining the multi-line structure.

examples/blocksparse_attention/example_tilelang_sparse_gqa_decode_varlen_indice.py (2)

196-196: LGTM! Modern super() syntax.

The simplified super().__init__() call is the recommended Python 3 idiom, reducing boilerplate and potential errors.


372-374: LGTM! Clean multi-line f-string formatting.

The conversion from .format() to f-strings improves readability and performance while maintaining the multi-line structure.

examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py (2)

3-3: LGTM! Essential for modern type hints.

The from __future__ import annotations import enables the int | None syntax used throughout the file (lines 209, 338, 383) without requiring typing.Union. This is the recommended approach for Python 3.10+ style type hints in earlier Python versions.


438-446: LGTM! Clean f-string migration.

The conversion to f-strings improves readability and follows modern Python best practices. All format specifiers are correctly preserved (.2f for floating-point precision).

examples/flash_attention/example_mha_bwd.py (2)

263-263: LGTM! Generator expression for unpacking is idiomatic.

The change from list comprehension to generator expression aligns with pyupgrade's UP027 rule. For unpacking into a fixed number of variables, generators are preferred as they avoid creating an intermediate list.


339-343: LGTM! F-strings improve readability.

The conversion from .format() to f-strings follows pyupgrade's UP031/UP032 rules. The formatting specifications are correctly preserved, and f-strings offer better readability and performance.

examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py (2)

3-3: LGTM!

The from __future__ import annotations import is correctly placed and enables the use of modern type hint syntax (e.g., int | None, tuple[...] at lines 202, 331, 358) while maintaining compatibility with Python 3.7-3.9.


423-427: LGTM!

The conversion from .format() to f-strings is correct and maintains identical formatting. The changes improve readability and align with modern Python best practices.

Note: The AI summary mentions a speedup calculation line that is not present in the annotated code changes. This may indicate the summary is inconsistent or referencing changes outside the shown diff.

examples/linear_attention/example_mamba_chunk_scan.py (1)

232-236: LGTM! Clean modernization to f-strings.

The conversion from .format() to f-strings improves readability while preserving identical formatting behavior. This aligns well with the PR's objective to enable pyupgrade linting.

examples/attention_sink/example_mha_sink_bwd_bhsd.py (3)

2-2: LGTM: Modern annotation handling enabled.

The from __future__ import annotations import enables postponed evaluation of type hints (PEP 563), which allows the modern union syntax int | None used in function signatures (lines 395, 442) to work correctly across Python versions.


369-369: LGTM: Efficient unpacking with generator expression.

The conversion from list comprehension to generator expression is a minor memory optimization. During unpacking, the generator is fully consumed, so the behavior remains identical.


492-496: LGTM: Modern f-string formatting.

The conversion to f-strings improves readability while maintaining identical output formatting. All format specifiers (.2f) are correctly preserved.

examples/attention_sink/example_mha_sink_fwd_bhsd.py (2)

2-2: LGTM! Essential import for modern type hints.

The from __future__ import annotations import enables PEP 563 postponed evaluation, which is necessary for the modern union syntax (int | None) used throughout this file (lines 190, 246). This is a standard modernization pattern recommended by pyupgrade.


289-293: LGTM! Clean f-string conversions.

The migration from .format() to f-strings preserves the exact formatting semantics (2 decimal places) while improving readability and performance. The expression total_flops / latency * 1e-9 is correctly embedded in the f-strings without any functional changes.

tilelang/carver/roller/rasterization.py (1)

92-92: LGTM! Clean modernization to f-string.

The conversion from .format() to f-string is correct and aligns with the PR's pyupgrade objective. The variable panel_width is guaranteed to be defined at this point (handled by the None-check on lines 88-89), and the generated CUDA code maintains proper semicolon syntax.

tilelang/language/proxy.py (5)

4-4: LGTM! Clean import modernization.

Correctly removed Optional, Tuple, and Union from typing imports since they're replaced with Python 3.10+ built-in syntax (tuple, X | Y) throughout the file.


146-146: LGTM! Type hint modernization.

Correctly updated from Tuple[Any] to tuple[Any] per Python 3.10+ standards.


154-154: LGTM! Enhanced type safety.

The added type annotation tuple[Any] | PrimExpr | int correctly captures the method's overloaded parameter using modern union syntax, aligning with the runtime check at lines 158-159.


175-176: LGTM! Consistent type hint modernization.

Both shape and strides parameters correctly updated to use tuple[Any].


276-276: LGTM! Modern union syntax.

Correctly updated from Optional[str] to str | None, using the modern union operator. The from __future__ import annotations at line 3 ensures runtime compatibility while enabling this modern syntax.

examples/flash_attention/example_mha_fwd_bhsd.py (1)

206-210: LGTM! Clean f-string modernization.

The conversion from .format() to f-strings improves readability and follows modern Python conventions. The output formatting remains identical.

examples/flash_attention/example_mha_bwd_bhsd.py (2)

264-264: LGTM! Valid generator expression unpacking.

Replacing the list comprehension with a generator expression is valid—Python consumes the generator during tuple unpacking. This is slightly more memory efficient while maintaining identical behavior.


342-346: LGTM! Clean f-string modernization.

The conversion from .format() to f-strings improves readability and follows modern Python conventions.

examples/flash_attention/example_mha_fwd_bshd.py (1)

193-197: LGTM! Clean f-string modernization.

The conversion from .format() to f-strings improves readability and follows modern Python conventions.

examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py (2)

404-404: LGTM! Valid generator expression unpacking.

Replacing the list comprehension with a generator expression is valid and slightly more memory efficient.


542-546: LGTM! Clean f-string modernization.

The conversion from .format() to f-strings improves readability and follows modern Python conventions.

examples/flash_attention/example_mha_bwd_wgmma_pipelined.py (2)

273-273: LGTM! Valid generator expression unpacking.

Replacing the list comprehension with a generator expression is valid and slightly more memory efficient.


349-353: LGTM! Clean f-string modernization.

The conversion from .format() to f-strings improves readability and follows modern Python conventions.

examples/flash_attention/example_gqa_fwd_bshd.py (1)

251-255: LGTM! Clean f-string modernization.

The conversion from .format() to f-strings improves readability and follows modern Python conventions.

examples/flash_attention/example_mha_fwd_bshd_wgmma_pipelined.py (1)

198-202: LGTM! Clean f-string modernization.

The conversion from .format() to f-strings improves readability and follows modern Python conventions.

examples/flash_attention/example_mha_fwd_bhsd_wgmma_pipelined.py (1)

211-215: LGTM! Clean f-string modernization.

The conversion from .format() to f-strings improves readability and follows modern Python conventions. All changes in this PR consistently apply Python 3 modernization best practices.

examples/flash_attention/example_gqa_bwd.py (2)

382-382: LGTM: Memory-efficient unpacking.

The generator expression is more efficient than a list comprehension for unpacking, eliminating the intermediate list object.


520-524: LGTM: Modern string formatting.

The f-string conversions improve readability and performance compared to .format() style.

examples/linear_attention/example_mamba_chunk_state.py (2)

171-175: LGTM! Clean f-string conversion.

The migration from .format() to f-strings is correct and improves readability. These changes align with the PR's objective of enabling the pyupgrade linter in ruff.


43-43: Approve parentheses removal in decay_states
This change only removes redundant parentheses around the subtraction; functionality is unchanged.

tilelang/jit/adapter/utils.py (2)

1-4: LGTM! Import cleanup aligns with modern Python.

The removal of Dict from typing imports is correct, as the built-in dict generic can be used directly with Python 3.9+.


110-110: LGTM! Type annotation modernization is correct.

The signature update correctly uses modern Python type hint syntax. The Python version verification requested for get_annotated_mod (lines 67-72) also applies to this change.

tilelang/jit/kernel.py (1)

20-20: LGTM! Modernized class declaration.

Removed explicit object base class, which is implicit in Python 3+. This aligns with the pyupgrade rule UP004.

tilelang/intrinsics/wgmma_macro_generator.py (1)

408-408: LGTM! Modernized string formatting.

Converted to f-string for better readability and performance. This aligns with the pyupgrade rule UP032.

tilelang/carver/arch/metal.py (1)

1-1: LGTM! Enabled postponed annotation evaluation.

Added from __future__ import annotations to enable PEP 563, which supports the modern union type syntax (Target | str on line 12) and improves performance by deferring annotation evaluation.

examples/bitnet-1.58b/modeling_bitnet.py (1)

1515-1515: LGTM! Removed redundant file mode.

The default mode for open() is 'r' (read text), so explicitly specifying it is unnecessary. This aligns with the pyupgrade rule UP015.

tilelang/jit/adapter/cython/adapter.py (3)

105-105: LGTM! Removed redundant file mode.

The default mode for open() is 'r' (read text), so explicitly specifying it is unnecessary. This aligns with the pyupgrade rule UP015.


118-118: LGTM! Removed redundant file mode.

The default mode for open() is 'r' (read text), so explicitly specifying it is unnecessary. This aligns with the pyupgrade rule UP015.


135-135: LGTM! Removed redundant file mode.

The default mode for open() is 'r' (read text), so explicitly specifying it is unnecessary. This aligns with the pyupgrade rule UP015.

tilelang/carver/roller/bestfit.py (1)

20-20: LGTM! Modernized string formatting.

Converted to f-string for better readability and performance. This aligns with the pyupgrade rule UP032.

examples/bitnet-1.58b/maint/create_bitblas_ckpt.py (1)

88-88: LGTM! Removed redundant file mode.

The default mode for open() is 'r' (read text), so explicitly specifying it is unnecessary. This aligns with the pyupgrade rule UP015.

tilelang/version.py (1)

24-24: LGTM! Removed redundant file mode.

The default mode for open() is 'r' (read text), so explicitly specifying it is unnecessary. This aligns with the pyupgrade rule UP015.

tilelang/carver/arch/cuda.py (1)

94-94: LGTM! Python 3 modernization.

Removing the explicit object base class is correct. In Python 3, all classes implicitly inherit from object, making this explicit declaration redundant.

tilelang/carver/roller/hint.py (1)

150-150: LGTM! Python 3 modernization.

The removal of the explicit object base class is appropriate for Python 3, where all classes implicitly inherit from object.

examples/hadamard_transform/example_hadamard.py (1)

154-154: LGTM! F-string modernization.

The conversion from .format() to f-string is a good modernization. F-strings are more readable and slightly more performant in Python 3.6+.

examples/bitnet-1.58b/vllm_workspace/conftest.py (1)

37-37: LGTM! Simplified file opening.

Removing the explicit "r" mode is appropriate since it's the default for open(). This makes the code cleaner without changing behavior.

tilelang/quantize/lop3.py (1)

1189-1191: LGTM! F-string modernization.

The conversion to f-strings improves readability. The error message and function name construction are now clearer and more idiomatic.

tilelang/jit/adapter/libgen.py (1)

32-32: LGTM! Python 3 modernization.

Removing the explicit object base class aligns with Python 3 best practices, where all classes implicitly inherit from object.

tilelang/cache/kernel_cache.py (1)

366-366: LGTM! Simplified file opening.

Omitting the explicit "r" mode is appropriate since it's the default for open(). This simplifies the code without affecting behavior.

tilelang/autotuner/param.py (3)

249-249: LGTM! Simplified file opening.

Removing the explicit "r" mode is correct since it's the default for open(). This makes the code cleaner.


318-318: LGTM! Simplified file opening.

Omitting the explicit "r" mode is appropriate and aligns with Python best practices.


330-330: LGTM! Simplified file opening.

The removal of the explicit "r" mode is correct and makes the code more concise.

tilelang/intrinsics/mma_macro_generator.py (2)

28-28: LGTM! Python 3 modernization.

Removing the explicit object base class is the recommended style for Python 3, as all classes implicitly inherit from object.


521-521: LGTM! Improved string formatting.

The f-string provides the same functionality as the previous .format() call with clearer syntax.

docs/conf.py (2)

4-4: LGTM! Modernized string formatting.

The f-string provides the same result as the previous % formatting with improved readability.


7-7: LGTM! Simplified file opening.

Removing the explicit "r" mode is appropriate since text read mode is the default for open().

examples/norm/rms_norm.py (1)

74-76: LGTM! Improved print formatting.

The f-string conversion maintains the same output format while improving code readability.

examples/attention_sink/example_gqa_sink_bwd_bhsd.py (3)

2-2: LGTM! Enables modern type hint syntax.

Adding from __future__ import annotations enables PEP 563 postponed evaluation of annotations, allowing the use of modern type hint syntax like int | None on Line 389.


361-361: LGTM! Memory-efficient unpacking.

The generator expression is more memory-efficient than the list comprehension since values are consumed immediately during unpacking. The behavior is identical for this use case.


488-492: LGTM! Improved string formatting.

The f-string conversions maintain the same output format with better readability.

tilelang/intrinsics/mfma_macro_generator.py (1)

14-14: LGTM! Python 3 modernization.

Removing the explicit object base class aligns with Python 3 conventions, as all classes implicitly inherit from object.

tilelang/carver/roller/shape_inference/common.py (1)

7-7: LGTM! Python 3 modernization.

Removing explicit object base classes aligns with Python 3 conventions, as all classes implicitly inherit from object.

Also applies to: 21-21

examples/amd/example_amd_flash_attn_bwd.py (2)

247-247: LGTM! Memory-efficient unpacking.

The generator expression is more memory-efficient than a list comprehension since values are consumed immediately during unpacking. The behavior is identical for this use case.


346-350: LGTM! Improved string formatting.

The f-string conversions maintain the same output format with better readability and performance.

tilelang/tileop/gemm/gemm_base.py (1)

11-12: Ensure subclasses of GemmBase remain instantiable after removing @dataclass.

Removing @dataclass drops the generated __init__, so any subclass without its own constructor will inherit object.__init__ and break when passed gemm_node. We found no direct instantiations of GemmBase or its subclasses (tilelang/tileop/gemm/gemm_wgmma.py, tilelang/tileop/gemm/gemm_mma.py, tilelang/primitives/gemm/gemm_mma.py), but none define __init__. Confirm these classes aren’t instantiated at runtime or add explicit __init__ methods.

setup.py (1)

140-140: LGTM!

The removal of explicit "r" mode from open() calls is a safe modernization since text mode is the default in Python 3. These changes align with enabling the pyupgrade linter.

Also applies to: 171-171, 711-712, 723-724, 740-741

tilelang/carver/roller/node.py (2)

32-32: LGTM!

Removing explicit object base classes is correct for Python 3, where all classes implicitly inherit from object.

Also applies to: 93-93


304-304: Verify that instance lifetime justifies method-level caching.

The static analysis tool (B019) flags that using @functools.lru_cache on methods can prevent garbage collection of instances. The cache holds references to self, which can lead to memory leaks if instances are meant to be short-lived.

If PrimFuncNode instances have long lifetimes or are reused frequently, this is acceptable. Otherwise, consider using a separate cache or @lru_cache with maxsize to limit memory growth.

Also applies to: 421-421

tilelang/carver/roller/shape_inference/tir.py (1)

50-50: LGTM!

Removing explicit object base classes is correct for Python 3.

Also applies to: 79-79

tilelang/jit/adapter/wrapper.py (4)

179-179: LGTM!

Removing explicit object base classes is correct for Python 3.

Also applies to: 895-895, 1062-1062


325-326: LGTM!

The f-string conversion for CUDA kernel launch configuration is correct and more readable.


350-351: LGTM!

The f-string conversion for kernel launch code is correct and improves readability.


993-993: LGTM!

The f-string conversion is correct.

Comment on lines 67 to +72
def get_annotated_mod(
func_or_mod: Union[tir.PrimFunc, tvm.IRModule],
target: Union[str, Target] = "auto",
target_host: Optional[Union[str, Target]] = None,
func_or_mod: tir.PrimFunc | tvm.IRModule,
target: str | Target = "auto",
target_host: str | Target | None = None,
model_type: Literal["device", "host", "all"] = "all",
) -> Union[IRModule, tuple[IRModule, IRModule]]:
) -> IRModule | tuple[IRModule, IRModule]:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

Verify Python version requirement for PEP 604 union syntax.

The type annotation updates look correct and align with modern Python best practices. However, the | union operator requires Python 3.10+.

Run the following script to verify the project's minimum Python version:


🏁 Script executed:

#!/bin/bash
# Description: Check the project's minimum Python version requirement

# Check pyproject.toml
if [ -f pyproject.toml ]; then
  echo "=== pyproject.toml ==="
  rg -A 3 "requires-python|python_requires" pyproject.toml
fi

# Check setup.py
if [ -f setup.py ]; then
  echo "=== setup.py ==="
  rg -A 3 "python_requires" setup.py
fi

# Check setup.cfg
if [ -f setup.cfg ]; then
  echo "=== setup.cfg ==="
  rg -A 3 "python_requires" setup.cfg
fi

# Check CI configuration for Python versions
echo "=== CI Python versions ==="
fd -e yml -e yaml . .github | xargs rg -l "python-version" | xargs cat

Length of output: 15112


Enforce Python ≥3.10 for PEP 604 union syntax
The new signature uses | unions (PEP 604), which require Python 3.10+. Update python_requires to >=3.10 in setup.py (and pyproject.toml if used), or revert to typing.Union/Optional for compatibility with the current >=3.8.

🤖 Prompt for AI Agents
In tilelang/jit/adapter/utils.py around lines 67 to 72 the function signature
uses PEP 604 union syntax (e.g. A | B) which requires Python >=3.10; either
update packaging to require Python 3.10+ by setting python_requires=">=3.10" in
setup.py and updating pyproject.toml if present (and adjust CI/test matrices) or
revert the signature to use typing.Union and typing.Optional (e.g.
Union[tir.PrimFunc, tvm.IRModule], Union[str, Target], Optional[Target]) and
ensure necessary typing imports are added/updated.

@oraluben oraluben changed the title Enable pyupgrade linter in ruff [Lint] Enable pyupgrade linter in ruff Oct 10, 2025
@LeiWang1999
Copy link
Member

@XuehaiPan Would you mind help review this pr?

LeiWang1999
LeiWang1999 previously approved these changes Oct 10, 2025
Copy link
Contributor

@XuehaiPan XuehaiPan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall LGTM.

Some small nits:

  1. Add __future__ imports only when needed.

  2. We'd best add the tool.ruff.target-version to pyproject.toml when UP code is enabled.

    [tool.ruff]
    target-version = "py38"

@oraluben
Copy link
Contributor Author

oraluben commented Oct 10, 2025

Thanks! @XuehaiPan

  1. Add __future__ imports only when needed.

This should be covered by the added FA102 rule.

2. We'd best add the tool.ruff.target-version to pyproject.toml when UP code is enabled.

I steal a commit from your PR #955 to resolve that.

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 22

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (19)
tilelang/language/persistent.py (1)

18-18: Update docstring to match the modernized type annotation.

The docstring still references List[tir.PrimExpr], but the function signature now uses list[tir.PrimExpr]. Update the docstring for consistency.

Apply this diff:

-    domain : List[tir.PrimExpr]
+    domain : list[tir.PrimExpr]
tilelang/primitives/gemm/base.py (1)

229-239: Update the docstring to match the new type annotation syntax.

The method parameter type was correctly updated to int | None, but the docstring at line 239 still references Optional[int]. Please update the docstring to maintain consistency.

Apply this diff to update the docstring:

         Parameters
         ----------
-        threads : Optional[int]
+        threads : int | None
             The total number of threads in a block. Must be provided
             if any block partition parameter is not already set.
tilelang/language/frame.py (1)

201-210: Update docstring to match the new type annotation syntax.

The return type annotation on line 201 uses the modern PrimExpr | None syntax, but the docstring on line 208 still references Optional[PrimExpr]. For consistency, update the docstring to use the new union syntax.

Apply this diff to update the docstring:

     """Get the value bound to a variable in the current let frame stack.

     Args:
         var (Var): The variable to look up

     Returns:
-        Optional[PrimExpr]: The bound value if found, None otherwise
+        PrimExpr | None: The bound value if found, None otherwise
     """
examples/bitnet-1.58b/vllm_workspace/conftest.py (1)

220-232: Restore the explicit None check for images.

Switching the guard to if images: makes an empty list skip the length assertion, yet the code still does images[i] later when images is not None, raising an IndexError. Revert to the explicit None check to keep the old, safe behavior.

-        if images:
+        if images is not None:
             assert len(prompts) == len(images)
tilelang/language/tir/ir.py (2)

10-32: Fix implicit Optional violations and update docstring.

Three issues need attention:

  1. Line 11: stop: PrimExpr = None violates PEP 484 (implicit Optional). Should be PrimExpr | None = None.
  2. Line 13: annotations: dict[str, Any] = None violates PEP 484. Should be dict[str, Any] | None = None.
  3. Line 24: Docstring still references Dict[str, Any] instead of dict[str, Any].

Apply this diff:

-def serial(start: PrimExpr,
-           stop: PrimExpr = None,
+def serial(start: PrimExpr, 
+           stop: PrimExpr | None = None,
            *,
-           annotations: dict[str, Any] = None) -> frame.ForFrame:
+           annotations: dict[str, Any] | None = None) -> frame.ForFrame:
     """The serial For statement.

     Parameters
     ----------
     start : PrimExpr
         The minimum value of iteration.

     stop : PrimExpr
         The maximum value of iteration.

-    annotations : Dict[str, Any]
+    annotations : dict[str, Any] | None
         The optional annotations of the For statement.

35-57: Apply the same fixes to remaining functions.

The parallel, vectorized, unroll, and thread_binding functions have the same three issues as serial:

  1. stop: PrimExpr = Nonestop: PrimExpr | None = None
  2. annotations: dict[str, Any] = Noneannotations: dict[str, Any] | None = None
  3. Docstrings reference Dict[str, Any] → update to dict[str, Any] | None

For parallel (lines 35-57):

-def parallel(start: PrimExpr,
-             stop: PrimExpr = None,
+def parallel(start: PrimExpr, 
+             stop: PrimExpr | None = None,
              *,
-             annotations: dict[str, Any] = None) -> frame.ForFrame:
+             annotations: dict[str, Any] | None = None) -> frame.ForFrame:
     """The parallel For statement.

     Parameters
     ----------
     start : PrimExpr
         The minimum value of iteration.

     stop : PrimExpr
         The maximum value of iteration.

-    annotations : Dict[str, Any]
+    annotations : dict[str, Any] | None
         The optional annotations of the For statement.

For vectorized (lines 60-82):

-def vectorized(start: PrimExpr,
-               stop: PrimExpr = None,
+def vectorized(start: PrimExpr, 
+               stop: PrimExpr | None = None,
                *,
-               annotations: dict[str, Any] = None) -> frame.ForFrame:
+               annotations: dict[str, Any] | None = None) -> frame.ForFrame:
     """The vectorized For statement.

     Parameters
     ----------
     start : PrimExpr
         The minimum value of iteration.

     stop : PrimExpr
         The maximum value of iteration.

-    annotations : Dict[str, Any]
+    annotations : dict[str, Any] | None
         The optional annotations of the For statement.

For unroll (lines 85-107):

-def unroll(start: PrimExpr,
-           stop: PrimExpr = None,
+def unroll(start: PrimExpr, 
+           stop: PrimExpr | None = None,
            *,
-           annotations: dict[str, Any] = None) -> frame.ForFrame:
+           annotations: dict[str, Any] | None = None) -> frame.ForFrame:
     """The unrolled For statement.

     Parameters
     ----------
     start : PrimExpr
         The minimum value of iteration.

     stop : PrimExpr
         The maximum value of iteration.

-    annotations : Dict[str, Any]
+    annotations : dict[str, Any] | None
         The optional annotations of the For statement.

For thread_binding (lines 110-138):

 def thread_binding(
     start: PrimExpr,
-    stop: PrimExpr = None,
-    thread: str = None,
+    stop: PrimExpr | None = None,
+    thread: str | None = None,
     *,
-    annotations: dict[str, Any] = None,
+    annotations: dict[str, Any] | None = None,
 ) -> frame.ForFrame:
     """The thread-binding For statement.

     Parameters
     ----------
     start : PrimExpr
         The minimum value of iteration.

     stop : PrimExpr
         The maximum value of iteration.

     thread : str
         The thread for loop variable to bind.

-    annotations : Dict[str, Any]
+    annotations : dict[str, Any] | None
         The optional annotations of the For statement.

Note: thread_binding also has thread: str = None that needs fixing.

Also applies to: 60-82, 85-107, 110-138

examples/fusedmoe/example_fusedmoe_torch.py (1)

10-16: Don't use PEP 604/585 syntax under a Python 3.8 target.

The new annotations (int | None, tuple[...], bare dict) rely on PEP 604/585 features that aren’t available on Python 3.8. Even with from __future__ import annotations, any runtime evaluation of these hints (e.g. typing.get_type_hints) will raise, breaking the stated goal of keeping 3.8 compatibility. Please revert to the typing.Optional, typing.Tuple, and typing.Dict forms (or equivalent) in this module before enabling the lint.

Apply this diff:

-from __future__ import annotations
+from __future__ import annotations
+from typing import Dict, Optional, Tuple
@@
-    def __init__(self, config: dict, d_expert: int | None = None):
+    def __init__(self, config: Dict, d_expert: Optional[int] = None):
@@
-    def __init__(self, config: dict):
+    def __init__(self, config: Dict):
@@
-    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
@@
-    def __init__(self, config: dict):
+    def __init__(self, config: Dict):
@@
-def ref_kernel(data: tuple[torch.Tensor, dict, dict]) -> torch.Tensor:
+def ref_kernel(data: Tuple[torch.Tensor, Dict, Dict]) -> torch.Tensor:
@@
-                   seed: int) -> tuple[torch.Tensor, dict, dict]:
+                   seed: int) -> Tuple[torch.Tensor, Dict, Dict]:

Also applies to: 37-38, 47-48, 100-146

tilelang/language/experimental/gemm_sp.py (1)

45-56: Type hint modernized. Update docstring and consider adding return type annotation.

The PEP 604 union syntax is correctly applied to the parameter type hint.

However, the docstrings (lines 49, 52) still reference the old Union[tir.Buffer, tir.Var] syntax. Additionally, consider adding a return type annotation for completeness.

Apply this diff to update the docstrings and add the return type annotation:

-    def legalize_arguments(arg: tir.Buffer | tir.Var):
+    def legalize_arguments(arg: tir.Buffer | tir.Var) -> tir.Buffer | tir.Var:
         """Convert let-bound variables to their corresponding buffers.
 
         Args:
-            arg (Union[tir.Buffer, tir.Var]): Input argument to legalize
+            arg (tir.Buffer | tir.Var): Input argument to legalize
 
         Returns:
-            Union[tir.Buffer, tir.Var]: The legalized argument
+            tir.Buffer | tir.Var: The legalized argument
         """
tilelang/carver/matmul_analysis.py (1)

337-344: Correct the get_ordered_axes return annotation

The helper still builds and returns a list, but the annotation now advertises set[Var]. This will confuse type-checkers (and readers) because downstream code indexes it (axes[-1]). Please change the return type back to list[Var] here and in the mirrored helper inside analysis_tensorcore_tags.

-    def get_ordered_axes(region: list[Range]) -> set[Var]:
+    def get_ordered_axes(region: list[Range]) -> list[Var]:
         axes: list[Var] = []
tilelang/language/kernel.py (1)

195-200: threads property includes the block frame; should return only threadIdx.{x,y,z}.

This currently returns 4 items (includes the last block frame). Align with get_thread_bindings and the doc intent.

     def threads(self) -> list[Var]:
         """
         Returns the thread indices from the topmost frame.
         """
-        return [frame.iter_var.var for frame in self.frames[-4:]]
+        # Exclude the trailing block frame; only return threadIdx.{x,y,z}
+        return [frame.iter_var.var for frame in self.frames[-4:-1]]
tilelang/engine/phase.py (1)

22-30: Guard None target before calling have_tma.

have_tma(None) will raise (accesses target.kind). Add a fast‑path for None.

 def allow_tma_and_warp_specialized(pass_ctx: PassContext | None = None,
                                    target: Target | None = None) -> bool:
     if pass_ctx is None:
         pass_ctx = tilelang.transform.get_pass_context()
+    if target is None:
+        return False
     if not have_tma(target):
         return False
     disable_tma_lower = pass_ctx.config.get("tl.disable_tma_lower", False)
     return not disable_tma_lower and allow_warp_specialized(pass_ctx=pass_ctx, target=target)

Reference: have_tma expects target.kind.name (see tilelang/contrib/nvcc.py). [Based on relevant code snippet]

tilelang/carver/arch/arch_base.py (1)

17-34: Duplicate attribute definitions overwrite initial values.

The attributes transaction_size and bandwidth are defined twice in the __init__ method:

  • First at lines 17-25 with comments describing their purpose
  • Again at lines 32-34 with comments about units

The second definitions (lines 32, 34) overwrite the first ones (lines 17, 22), making the initial assignments redundant. This appears to be a pre-existing issue that should be addressed.

Remove the duplicate definitions and consolidate into a single set:

         self.max_smem_usage: int = 0  # The maximum shared memory usage allowed
-        self.bandwidth: list[int] = [
-            0,
-            0,
-        ]  # Bandwidth specifications, possibly including peak and sustained rates
         self.platform: str = "unknown"  # The platform or manufacturer of the device
         self.compute_capability: str = (
             "unknown"  # The compute capability, indicating the feature set and performance level
         )
         self.l2_cache_size_bytes: int = 0
-        # the number of transaction size in bytes
-        self.transaction_size: list[int] = [0, 0]  # in bytes
-        # bandwidth in MB/s, will be used for recommend basic tile size
-        self.bandwidth: list[int] = [0, 0]
+        # Consolidate comments: The number of transaction size in bytes
+        self.transaction_size: list[int] = [0, 0]
+        # Consolidate comments: Bandwidth in MB/s, will be used for recommend basic tile size (possibly including peak and sustained rates)
+        self.bandwidth: list[int] = [0, 0]

Alternatively, if the second definitions should replace the first, remove lines 17-25:

         self.sm_partition: int = 0  # The number of streaming multiprocessor partitions
-        self.transaction_size: list[int] = [
-            0,
-            0,
-        ]  # The size of memory transactions, typically in bytes
-        self.max_smem_usage: int = 0  # The maximum shared memory usage allowed
-        self.bandwidth: list[int] = [
-            0,
-            0,
-        ]  # Bandwidth specifications, possibly including peak and sustained rates
+        self.max_smem_usage: int = 0  # The maximum shared memory usage allowed
         self.platform: str = "unknown"  # The platform or manufacturer of the device
tilelang/engine/lower.py (1)

128-133: Remove unused target parameter from canon_target_host.

The target argument isn’t referenced; drop it from the signature and update both call sites.

--- a/tilelang/engine/lower.py
@@ -128,7 +128,6 @@ def canon_target_host(
- def canon_target_host(target: str | Target, target_host: str | Target | None):
+ def canon_target_host(target_host: str | Target | None):
     if not target_host:
         target_host = "llvm" if tvm.runtime.enabled("llvm") else "c"
     return target_host
@@ -216,7 +215,7 @@ def lower(...):
-    target_host = canon_target_host(target, target_host)
+    target_host = canon_target_host(target_host)
--- a/tilelang/jit/adapter/utils.py
@@ -85,7 +85,7 @@ def some_util(...):
-    target_host = tvm.target.Target.canon_target(canon_target_host(target, target_host))
+    target_host = tvm.target.Target.canon_target(canon_target_host(target_host))
tilelang/carver/arch/driver/cuda_driver.py (1)

124-143: Don’t silently return None from get_device_attribute

The new try/except now returns None on failure while the signature remains -> int. Existing callers expect an int (or a raised error) and will now receive None, causing downstream breakage or type confusion. Please restore the previous behavior—let the RuntimeError propagate (or at least re-raise it) and keep returning an int, or explicitly change the return type and update every caller to handle the None case.

tilelang/jit/adapter/nvrtc/adapter.py (1)

163-167: Guard buffer_map lookup before indexing

PrimFunc.params often include scalar vars in addition to buffers. Those scalars are absent from buffer_map, so the new direct lookup raises a KeyError during adapter initialisation. Please skip params that aren’t backed by buffers, like the ctypes adapter still does.

Apply:

-        for i, param in enumerate(params):
-            buffer = buffer_map[param]
+        for i, param in enumerate(params):
+            if param not in buffer_map:
+                continue
+            buffer = buffer_map[param]
tilelang/jit/adapter/ctypes/adapter.py (2)

226-233: Unpack bug after changing dynamic_symbolic_map to (id, i, j)

dynamic_symbolic_map now stores triples (id, buffer_idx, dim). This two-value unpack will raise “too many values to unpack”.

-                    if isinstance(s, tir.Var):
-                        ref_tensor_idx, ref_shape_idx = self.dynamic_symbolic_map[s]
-                        shape.append(ins[ref_tensor_idx].shape[ref_shape_idx])
+                    if isinstance(s, tir.Var):
+                        _, ref_tensor_idx, ref_shape_idx = self.dynamic_symbolic_map[s]
+                        shape.append(ins[ref_tensor_idx].shape[ref_shape_idx])

239-245: Use args, not ins, for dynamic symbolics
When appending dynamic dim/stride values, index into args (which includes both inputs and outputs) rather than ins (inputs-only) to support symbols on output buffers. Replace ins[buffer_idx] with args[buffer_idx] in the loop at lines 239–245.

tilelang/autotuner/tuner.py (1)

232-256: generate_cache_key return type is wrong

Function returns a hex string, not AutotuneResult | None.

-    def generate_cache_key(self, parameters: dict[str, Any]) -> AutotuneResult | None:
+    def generate_cache_key(self, parameters: dict[str, Any]) -> str:
tilelang/jit/adapter/wrapper.py (1)

493-499: Use of undefined variable 'function_name'

Inside host_mod loop, function_name is not defined; this will raise at runtime. l2_persistent_map likely applies to all kernels; set the map directly.

-            if "l2_persistent_map" in func.attrs:
-                self.l2_persistent_map[function_name] = func.attrs["l2_persistent_map"]
+            if "l2_persistent_map" in func.attrs:
+                self.l2_persistent_map = func.attrs["l2_persistent_map"]
♻️ Duplicate comments (2)
setup.py (1)

312-312: Remove unused noqa (SIM115 not enabled).

Ruff flags this as an unused suppression.

-        return open(get_path("README.md"), encoding="utf-8").read()  # noqa: SIM115
+        return open(get_path("README.md"), encoding="utf-8").read()
tilelang/carver/roller/node.py (1)

305-307: Restore lru_cache() invocation

Dropping the parentheses passes the method object in as maxsize, so definition now raises TypeError: 'function' object cannot be interpreted as an integer. Please revert to calling the decorator (applies to both cached methods).

Fix:

-    @functools.lru_cache
+    @functools.lru_cache()
@@
-    @functools.lru_cache
+    @functools.lru_cache()

Also applies to: 420-422

🧹 Nitpick comments (23)
examples/dequantize_gemm/example_dequant_gemv_fp16xint4.py (1)

76-76: Remove redundant self-assignment.

Line 76 assigns import_source to itself, which has no effect. This line can be safely removed.

Apply this diff to remove the redundant line:

     assert import_source is not None, "lop3_intrin_info is not found"
     assert func_name is not None, "lop3_intrin_info is not found"
-    import_source = import_source
tilelang/utils/language.py (1)

88-98: Update docstring to match the new type annotation.

The function signature correctly uses the modern list[int] syntax, but the docstring at line 93 still references List[int]. For consistency, update the docstring to match the signature.

Apply this diff to update the docstring:

     Args:
-        array (List[int]): The array of integers to reduce.
+        array (list[int]): The array of integers to reduce.
 
     Returns:
testing/python/kernel/test_tilelang_kernel_fp8_gemv_simt.py (1)

25-26: Type annotation modernization looks good.

The migration from Optional[int] to int | None syntax is correct and aligns with the PR's goal of enabling pyupgrade linter rules. The runtime behavior is unchanged since the defaults and assertions remain the same.

Optional refactor: Consider tightening the type hints.

Since both parameters have non-None defaults (4 and 32) and are immediately validated as non-None (lines 28-31), the type hints could be simplified to just int rather than int | None:

-    n_partition: int | None = 4,
-    reduce_thread: int | None = 32,
+    n_partition: int = 4,
+    reduce_thread: int = 32,

This would make the signature more accurate—the function doesn't meaningfully accept None since it's rejected immediately. However, this is a pre-existing pattern from the original Optional[int] annotations and is outside the scope of this linting PR.

tilelang/language/warpgroup.py (1)

48-50: Consider simplifying the list construction.

The current loop-based approach could be streamlined using a list constructor, though this is beyond the scope of the linting changes.

Apply this diff to simplify:

-    warp_group_ids: list[int] = []
-    for warp_group_id in warp_group_idx:
-        warp_group_ids.append(warp_group_id)
+    warp_group_ids: list[int] = list(warp_group_idx)
examples/bitnet-1.58b/vllm_workspace/utils.py (1)

1-2: Clean up unused imports and modernize type aliases for consistency.

After updating the function signatures to use built-in list, the List import from typing on line 2 is no longer used in the function signatures. Additionally, for consistency with the modernized function signatures, the type aliases TokensText (line 4) and TokensTextLogprobs (line 27) should also be updated to use built-in generics instead of Tuple and List from typing.

Apply this diff to modernize the type aliases and clean up imports:

-from typing import Dict, List, Tuple
+from typing import Dict


-TokensText = Tuple[List[int], str]
+TokensText = tuple[list[int], str]


-TokensTextLogprobs = Tuple[List[int], str, List[Dict[int, float]]]
+TokensTextLogprobs = tuple[list[int], str, list[Dict[int, float]]]

Note: If Dict is also unused elsewhere in the codebase, consider removing it too. The modernization to dict[int, float] would require postponed evaluation support, which is already enabled via the from __future__ import annotations statement.

tilelang/contrib/cc.py (1)

211-211: LGTM! Type annotation correctly modernized.

The change from typing.Dict[str, str] to dict[str, str] is correct and aligns with modern Python typing practices (PEP 585). The future annotations import at line 18 ensures Python 3.8 compatibility.

Optional: Consider updating the docstring for consistency.

The docstring at line 224 still uses the old-style Dict[str, str] notation:

    Returns
    -------
    symbol_section_map: Dict[str, str]
        A map from defined global symbol to their sections

For consistency with the actual annotation, you could update it to:

     Returns
     -------
-    symbol_section_map: Dict[str, str]
+    symbol_section_map: dict[str, str]
         A map from defined global symbol to their sections
examples/fusedmoe/example_fusedmoe_tilelang.py (4)

274-278: LGTM! Type hints modernized.

The conversion to built-in generic types (dict, int | None) is correct and consistent with PEP 585/604 standards.

Optionally, consider making the dict type more specific for better type safety:

     def __init__(self,
-                 config: dict,
+                 config: dict[str, Any],
                  gate: torch.Tensor,
                  up: torch.Tensor,
                  down: torch.Tensor,
                  d_expert: int | None = None):

This would require importing Any from typing if not already imported.


298-298: LGTM! Type hints modernized.

The conversion to built-in dict type is correct and consistent with PEP 585 standards.

Optionally, consider making the dict types more specific:

-    def __init__(self, config: dict, weights: dict):
+    def __init__(self, config: dict[str, Any], weights: dict[str, torch.Tensor]):

317-320: LGTM! Type hints modernized.

The conversion to built-in dict types is correct and consistent with PEP 585 standards.

Optionally, consider making the dict types more specific:

     def __init__(self,
-                 config: dict,
+                 config: dict[str, Any],
                  shared_kernel: tilelang.JITKernel,
                  routed_kernel: tilelang.JITKernel,
-                 weights: dict,
+                 weights: dict[str, torch.Tensor],
                  padding_M: int = 128):

478-478: LGTM! Type hint modernized.

The conversion to built-in tuple type is correct and consistent with PEP 585 standards.

Optionally, consider making the dict types more specific to match the documented structure:

-def custom_kernel(data: tuple[torch.Tensor, dict, dict]) -> torch.Tensor:
+def custom_kernel(data: tuple[torch.Tensor, dict[str, torch.Tensor], dict[str, Any]]) -> torch.Tensor:
examples/cast/example_per_token_cast_to_fp8.py (1)

103-117: All print statements use f-strings; refactor Tuple imports

  • No remaining .format() or % formatting in print() calls.
  • Replace from typing import Tuple with built-in tuple[...] annotations in:
    • maint/precision/compare_ops.py
    • tilelang/language/ast/ir.py
    • examples/deepseek_v32/inference/kernel.py
    • examples/deepseek_v32/inference/model.py
    • examples/bitnet-1.58b/vllm_workspace/utils.py
    • examples/deepseek_v32/utils.py
tilelang/contrib/nvrtc.py (1)

16-17: Update docstrings to match modernized type annotations.

The parameter type annotations have been correctly modernized to use PEP 604 union syntax (int | None, str | list[str] | None). However, the docstrings at lines 29 and 32 still reference the old typing notation (Optional[int], Optional[Union[str, List[str]]]).

Consider updating the docstring to match the new annotation style for consistency:

-    arch : Optional[int]
+    arch : int | None
         The cuda architecture code.
 
-    options : Optional[Union[str, List[str]]]
+    options : str | list[str] | None
         The additional options.
tilelang/language/builtin.py (1)

173-173: Consider expanding type hints to match implementation.

The type hints for mbarrier/barrier_id specify int | PrimExpr | tir.Call, but the implementations (lines 212-219, 230-237) also accept tir.BufferLoad and tir.Buffer. Consider adding these types to the annotations for more accurate API documentation:

-def mbarrier_wait_parity(mbarrier: int | PrimExpr | tir.Call, parity: int | Var):
+def mbarrier_wait_parity(mbarrier: int | PrimExpr | tir.Call | tir.Buffer | tir.BufferLoad, parity: int | Var):

(Apply similar changes to mbarrier_arrive, barrier_wait, and barrier_arrive)

Also applies to: 223-223, 266-266, 281-281

setup.py (1)

124-129: Harden requirements parsing and set encoding.

Avoid env‑dependent defaults and stray entries in install_requires. Filter comments/empties and set UTF‑8.

-def get_requirements(file_path: str = "requirements.txt") -> list[str]:
-    """Get Python package dependencies from requirements.txt."""
-    with open(get_path(file_path)) as f:
-        requirements = f.read().strip().split("\n")
-    return requirements
+def get_requirements(file_path: str = "requirements.txt") -> list[str]:
+    """Get Python package dependencies from requirements.txt."""
+    with open(get_path(file_path), encoding="utf-8") as f:
+        requirements: list[str] = []
+        for line in f:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            requirements.append(line)
+    return requirements
tilelang/language/kernel.py (2)

210-214: Fix varargs annotation for blocks.

For varargs, the annotation applies to each element. Use tir.PrimExpr, not list[tir.PrimExpr].

-def Kernel(
-    *blocks: list[tir.PrimExpr],
+def Kernel(
+    *blocks: tir.PrimExpr,

158-164: Docstring doesn’t match return type.

Method returns a list of three bindings, not a single dim’s binding.

 def get_thread_bindings(self) -> list[Var]:
-    """
-    Returns the thread binding for the given dimension.
-    dim=0 corresponds to threadIdx.x, dim=1 to threadIdx.y, and dim=2 to threadIdx.z.
-    """
+    """
+    Returns all three thread bindings: threadIdx.x, threadIdx.y, threadIdx.z.
+    """
tilelang/language/copy.py (1)

11-16: Include tir.Var in type hints to match runtime handling.

Implementation accepts tir.Var (resolved via T.has_let_value/T.get_let_value). Reflect that in signatures.

-def copy(src: tir.Buffer | tir.BufferLoad | tir.BufferRegion,
-         dst: tir.Buffer | tir.BufferLoad,
+def copy(src: tir.Buffer | tir.BufferLoad | tir.BufferRegion | tir.Var,
+         dst: tir.Buffer | tir.BufferLoad | tir.Var,
          coalesced_width: int | None = None,
          disable_tma: bool = False,
          eviction_policy: Literal["evict_normal", "evict_first", "evict_last"] | None = None):
tilelang/carver/template/conv.py (1)

47-59: Restore the TileDevice parameter annotation

All other templates still surface the arch: TileDevice annotation, and BaseTemplate.get_hardware_aware_configs declares the same signature. Dropping it here hides the contract from static analysis and breaks consistency. Please keep the type annotation (re-importing TileDevice if necessary) so the override matches the base method.

-from ..roller import Hint
+from ..arch import TileDevice
+from ..roller import Hint
@@
-    def get_hardware_aware_configs(self, arch=None, topk=10) -> list[Hint]:
+    def get_hardware_aware_configs(self, arch: TileDevice = None, topk: int = 10) -> list[Hint]:
tilelang/carver/arch/cuda.py (1)

137-138: Fix the available_tensor_instructions annotation

The attribute is annotated as list[TensorInstruction] but initialized to None, which violates the declared type and will trip stricter type checkers now that we’ve enabled the new linting. Please declare it as list[TensorInstruction] | None (or initialize it with an empty list) to keep the annotation truthful.

tilelang/jit/adapter/ctypes/adapter.py (1)

102-106: Check library init result and surface error

For parity with the cython adapter and better diagnostics, check init() return and expose get_last_error.

-        self.lib = self.lib_generator.load_lib()
-        self.lib.init()
+        self.lib = self.lib_generator.load_lib()
+        self.lib.get_last_error.restype = ctypes.c_char_p
+        _res = self.lib.init()
+        if _res != 0:
+            _err = self.lib.get_last_error().decode("utf-8")
+            raise RuntimeError(f"Initialization failed: {_err}")
tilelang/jit/adapter/cython/adapter.py (3)

378-393: buffer_dtype_map key type annotation mismatch

Keys are buffer names (str), not tir.Var. Fix annotations.

-    def _process_buffer_dtype(self) -> dict[tir.Var, tuple[int, torch.dtype]]:
+    def _process_buffer_dtype(self) -> dict[str, tuple[int, torch.dtype]]:
@@
-        buffer_dtype_map = {}
+        buffer_dtype_map: dict[str, tuple[int, torch.dtype]] = {}

408-412: _process_static_buffer_infos return type annotations incorrect

Maps keyed by buffer.name (str) and static_contiguous_list holds (index, name).

-    def _process_static_buffer_infos(self) -> \
-            tuple[dict[tir.Var, tuple[int, list[tuple[int, int]]]],
-                  dict[tir.Var, tuple[int, list[tuple[int, int]]]],
-                  list[tuple[tir.Var]]]:
+    def _process_static_buffer_infos(self) -> \
+            tuple[dict[str, tuple[int, list[tuple[int, int]]]],
+                  dict[str, tuple[int, list[tuple[int, int]]]],
+                  list[tuple[int, str]]]:

442-467: buffer_device_map key type annotation mismatch

Same as others: key is buffer name (str).

-    def _process_buffer_device(self) -> dict[tir.Var, tuple[int, torch.device]]:
+    def _process_buffer_device(self) -> dict[str, tuple[int, torch.device]]:
@@
-        buffer_device_map = {}
+        buffer_device_map: dict[str, tuple[int, torch.device]] = {}

Comment on lines 7 to 10
def run_general_reduction_recommend_hints(structure: str = "SSR",
shape: List[int] = None,
shape: list[int] = None,
dtype: str = "float16",
topk: int = 20):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Fix implicit Optional in type annotations.

The shape parameter has a default value of None, but the type annotation list[int] doesn't include None as a valid type. PEP 484 prohibits implicit Optional. With from __future__ import annotations enabled, use the PEP 604 union syntax list[int] | None.

Apply this diff to fix both occurrences:

 def run_general_reduction_recommend_hints(structure: str = "SSR",
-                                          shape: list[int] = None,
+                                          shape: list[int] | None = None,
                                           dtype: str = "float16",
                                           topk: int = 20):
-def run_elementwise_recommend_hints(shape: list[int] = None,
+def run_elementwise_recommend_hints(shape: list[int] | None = None,
                                     dtype: str = "float16",
                                     topk: int = 20):

Also applies to: 31-33

🧰 Tools
🪛 Ruff (0.13.3)

8-8: PEP 484 prohibits implicit Optional

Convert to Optional[T]

(RUF013)

🤖 Prompt for AI Agents
In testing/python/carver/test_tilelang_carver_recommend_hints.py around lines
7-10 and 31-33, the function parameter `shape` is annotated as `list[int]` but
has a default of `None`, creating an implicit Optional; update both annotations
to use PEP 604 union syntax `list[int] | None` so the type explicitly allows
None (i.e., change `shape: list[int] = None` to `shape: list[int] | None = None`
in both locations).

Comment on lines +51 to 55
out_idx: list[int] | int | None = None
execution_backend: Literal["dlpack", "ctypes", "cython"] = "cython"
target: Literal['auto', 'cuda', 'hip'] = 'auto'
target_host: Union[str, Target] = None
target_host: str | Target = None
verbose: bool = False
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Include None in the annotation for target_host.

With pyupgrade’s RUF013 enabled, a default of None without None in the annotation fails lint. Please widen the type to keep the new lint clean.

-    target_host: str | Target = None
+    target_host: str | Target | None = None
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
out_idx: list[int] | int | None = None
execution_backend: Literal["dlpack", "ctypes", "cython"] = "cython"
target: Literal['auto', 'cuda', 'hip'] = 'auto'
target_host: Union[str, Target] = None
target_host: str | Target = None
verbose: bool = False
out_idx: list[int] | int | None = None
execution_backend: Literal["dlpack", "ctypes", "cython"] = "cython"
target: Literal['auto', 'cuda', 'hip'] = 'auto'
target_host: str | Target | None = None
verbose: bool = False
🤖 Prompt for AI Agents
In tilelang/autotuner/param.py around lines 51 to 55, the parameter annotation
for target_host is missing None even though its default is None; update the type
to include None (e.g. change the annotation to include | None or use Optional[])
so the declared type matches the default and satisfies the RUF013 lint rule.

Comment on lines +107 to 109
_kernel_parameters: tuple[str, ...] | None = None
_function_parameters: dict[str, Any] | None = None
_lock = threading.Lock() # For thread safety
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

_kernel_parameters/_function_parameters types regressed

_key is used as (args_tuple, sorted_kwargs_tuple). The current tuple[str, ...] is incorrect and will confuse tools.

-    _kernel_parameters: tuple[str, ...] | None = None
-    _function_parameters: dict[str, Any] | None = None
+    _kernel_parameters: tuple[tuple[Any, ...], tuple[tuple[str, Any], ...]] | None = None
+    _function_parameters: dict[str, Any] | None = None
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
_kernel_parameters: tuple[str, ...] | None = None
_function_parameters: dict[str, Any] | None = None
_lock = threading.Lock() # For thread safety
_kernel_parameters: tuple[tuple[Any, ...], tuple[tuple[str, Any], ...]] | None = None
_function_parameters: dict[str, Any] | None = None
_lock = threading.Lock() # For thread safety
🤖 Prompt for AI Agents
In tilelang/autotuner/tuner.py around lines 107 to 109, the type annotations for
_kernel_parameters and _function_parameters are incorrect—the code uses keys of
the form (args_tuple, sorted_kwargs_tuple), not tuple[str,...] or str. Change
both annotations to a dict whose keys are a tuple of (args tuple, sorted kwargs
tuple) — e.g. key type tuple[tuple[Any, ...], tuple[tuple[str, Any], ...]] — and
values as the appropriate parameter type (Any), keeping the optional None union
and leaving _lock as-is for thread safety.

Comment on lines +227 to 231
def set_kernel_parameters(self, k_parameters: tuple[str, ...], f_parameters: dict[str, Any]):
# for cache key generation
self._kernel_parameters = k_parameters
self._function_parameters = f_parameters

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

set_kernel_parameters signature should match actual payload

The first argument is that (args_tuple, kwargs_tuple) key, not a tuple[str, ...].

-    def set_kernel_parameters(self, k_parameters: tuple[str, ...], f_parameters: dict[str, Any]):
+    def set_kernel_parameters(self, k_parameters: tuple[tuple[Any, ...], tuple[tuple[str, Any], ...]], f_parameters: dict[str, Any]):
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
def set_kernel_parameters(self, k_parameters: tuple[str, ...], f_parameters: dict[str, Any]):
# for cache key generation
self._kernel_parameters = k_parameters
self._function_parameters = f_parameters
def set_kernel_parameters(
self,
k_parameters: tuple[tuple[Any, ...], tuple[tuple[str, Any], ...]],
f_parameters: dict[str, Any]
):
# for cache key generation
self._kernel_parameters = k_parameters
self._function_parameters = f_parameters
🤖 Prompt for AI Agents
In tilelang/autotuner/tuner.py around lines 227 to 231, the
set_kernel_parameters signature and typing are incorrect: the first argument is
the (args_tuple, kwargs_tuple) key, not a tuple[str, ...]. Change the parameter
type to reflect the actual payload (e.g., key: tuple[tuple[Any, ...], dict[str,
Any]] or a more specific alias) and update the parameter name if helpful (e.g.,
key or params_key); then assign the provided key and the f_parameters to the
internal fields (self._kernel_parameters = key; self._function_parameters =
f_parameters) and update any callers if necessary to match the new
signature/type.

def cached(
func: PrimFunc = None,
out_idx: List[int] = None,
out_idx: list[int] = None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Add explicit | None annotation.

Per PEP 484, parameters with None defaults must explicitly include None in the type annotation.

Apply this diff:

 def cached(
     func: PrimFunc = None,
-    out_idx: list[int] = None,
+    out_idx: list[int] | None = None,
     *args,
🧰 Tools
🪛 Ruff (0.13.3)

17-17: PEP 484 prohibits implicit Optional

Convert to Optional[T]

(RUF013)

🤖 Prompt for AI Agents
In tilelang/cache/__init__.py around line 17, the parameter annotation `out_idx:
list[int] = None` needs to explicitly include None per PEP 484; change the type
to `list[int] | None` (or `Optional[list[int]]` for Python <3.10) so the
signature becomes `out_idx: list[int] | None = None`.

Comment on lines 44 to 55
def __init__(
self,
func: PrimFunc = None,
out_idx: Union[List[int], int] = None,
out_idx: list[int] | int = None,
execution_backend: Literal["dlpack", "ctypes", "cython", "nvrtc"] = "cython",
target: Union[str, Target] = "auto",
target_host: Union[str, Target] = None,
target: str | Target = "auto",
target_host: str | Target = None,
verbose: bool = False,
pass_configs: Optional[Dict[str, Any]] = None,
pass_configs: dict[str, Any] | None = None,
from_database: bool = False,
compile_flags: Optional[List[str]] = None,
compile_flags: list[str] | None = None,
):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Fix implicit Optional annotations in JITKernel.__init__.

out_idx and target_host default to None, but their annotations omit it. RUF013 (pyupgrade) will flag this, blocking the new lint configuration. Please fold None into the unions.

-        out_idx: list[int] | int = None,
-        target: str | Target = "auto",
-        target_host: str | Target = None,
+        out_idx: list[int] | int | None = None,
+        target: str | Target = "auto",
+        target_host: str | Target | None = None,
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
def __init__(
self,
func: PrimFunc = None,
out_idx: Union[List[int], int] = None,
out_idx: list[int] | int = None,
execution_backend: Literal["dlpack", "ctypes", "cython", "nvrtc"] = "cython",
target: Union[str, Target] = "auto",
target_host: Union[str, Target] = None,
target: str | Target = "auto",
target_host: str | Target = None,
verbose: bool = False,
pass_configs: Optional[Dict[str, Any]] = None,
pass_configs: dict[str, Any] | None = None,
from_database: bool = False,
compile_flags: Optional[List[str]] = None,
compile_flags: list[str] | None = None,
):
def __init__(
self,
func: PrimFunc = None,
out_idx: list[int] | int | None = None,
execution_backend: Literal["dlpack", "ctypes", "cython", "nvrtc"] = "cython",
target: str | Target = "auto",
target_host: str | Target | None = None,
verbose: bool = False,
pass_configs: dict[str, Any] | None = None,
from_database: bool = False,
compile_flags: list[str] | None = None,
):
🧰 Tools
🪛 Ruff (0.13.3)

47-47: PEP 484 prohibits implicit Optional

Convert to Optional[T]

(RUF013)

🤖 Prompt for AI Agents
In tilelang/jit/kernel.py around lines 44 to 55, the parameters out_idx and
target_host default to None but their type annotations omit None; update the
signatures so None is included in the unions (e.g., out_idx: list[int] | int |
None and target_host: str | Target | None) to make the types explicit and
satisfy the pyupgrade RUF013 lint rule; ensure any other parameters with default
None are similarly annotated if applicable.

Comment on lines +10 to +13
A_sparse: tir.Buffer | tir.Var,
E: tir.Buffer | tir.Var,
B: tir.Buffer | tir.Var,
C: tir.Buffer | tir.Var,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Type hints correctly modernized. Update docstrings for consistency.

The PEP 604 union syntax (|) is correctly applied to the function parameters and works with the __future__ import for Python 3.8+ compatibility.

However, the docstrings (lines 27-30) still reference the old Union[tir.Buffer, tir.Var] syntax. Update them to match the actual type hints.

Apply this diff to update the docstrings:

     Args:
-        A_sparse (Union[tir.Buffer, tir.Var]): First input matrix dense values
-        E (Union[tir.Buffer, tir.Var]): First input matrix sparse metadata
-        B (Union[tir.Buffer, tir.Var]): Second input matrix
-        C (Union[tir.Buffer, tir.Var]): Output matrix for results
+        A_sparse (tir.Buffer | tir.Var): First input matrix dense values
+        E (tir.Buffer | tir.Var): First input matrix sparse metadata
+        B (tir.Buffer | tir.Var): Second input matrix
+        C (tir.Buffer | tir.Var): Output matrix for results
         transpose_A (bool, optional): Whether to transpose matrix A. Defaults to False.
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
A_sparse: tir.Buffer | tir.Var,
E: tir.Buffer | tir.Var,
B: tir.Buffer | tir.Var,
C: tir.Buffer | tir.Var,
"""
Args:
A_sparse (tir.Buffer | tir.Var): First input matrix dense values
E (tir.Buffer | tir.Var): First input matrix sparse metadata
B (tir.Buffer | tir.Var): Second input matrix
C (tir.Buffer | tir.Var): Output matrix for results
transpose_A (bool, optional): Whether to transpose matrix A. Defaults to False.
"""
🤖 Prompt for AI Agents
In tilelang/language/experimental/gemm_sp.py around lines 27 to 30, the
docstring parameter types still use the old Union[tir.Buffer, tir.Var] syntax;
update the docstring entries for A_sparse, E, B, and C to reflect the modern PEP
604 style used in the signature (tir.Buffer | tir.Var), keeping wording and
punctuation identical otherwise so the docs match the actual type hints.



def Parallel(*extents: tir.PrimExpr, coalesced_width: Optional[int] = None):
def Parallel(*extents: tir.PrimExpr, coalesced_width: int | None = None):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion | 🟠 Major

Update the docstring to match the new type annotation.

The parameter type annotation has been correctly modernized from Optional[int] to int | None, but the docstring at line 18 still references Optional[int]. Update the docstring for consistency.

Apply this diff to update the docstring:

     coalesced_width : Optional[int]
+    coalesced_width : int | None
         The coalesced width of the parallel loop.
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
def Parallel(*extents: tir.PrimExpr, coalesced_width: int | None = None):
coalesced_width : Optional[int]
coalesced_width : int | None
The coalesced width of the parallel loop.
🤖 Prompt for AI Agents
In tilelang/language/parallel.py around lines 9 to 18, the function signature
now uses the modern union type `int | None` for the parameter `coalesced_width`,
but the docstring still refers to `Optional[int]`; update the docstring text to
reference `int | None` (or simply "int or None") so it matches the annotation,
keeping the same description and formatting.

Comment on lines +1 to +12
from __future__ import annotations
import inspect
from typing import Callable, Optional, Union
from typing import Callable

import tvm.script.parser.tir.entry as _tir_entry
from tvm.tir.function import PrimFunc
from tvm.script.parser._core import parse, scan_macro, utils


def prim_func(func: Optional[Callable] = None,
def prim_func(func: Callable | None = None,
private: bool = False,
check_well_formed: bool = False) -> Union[PrimFunc, Callable]:
check_well_formed: bool = False) -> PrimFunc | Callable:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion | 🟠 Major

Update docstring to match new return type annotation.

The type annotations are correctly modernized using union syntax (Callable | None, PrimFunc | Callable). However, the docstring on line 31 still references the old Union[PrimFunc, Callable] syntax.

Apply this diff to update the docstring:

     Returns
     -------
-    res : Union[PrimFunc, Callable]
+    res : PrimFunc | Callable
         The parsed tir prim func.
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
from __future__ import annotations
import inspect
from typing import Callable, Optional, Union
from typing import Callable
import tvm.script.parser.tir.entry as _tir_entry
from tvm.tir.function import PrimFunc
from tvm.script.parser._core import parse, scan_macro, utils
def prim_func(func: Optional[Callable] = None,
def prim_func(func: Callable | None = None,
private: bool = False,
check_well_formed: bool = False) -> Union[PrimFunc, Callable]:
check_well_formed: bool = False) -> PrimFunc | Callable:
Returns
-------
res : PrimFunc | Callable
The parsed tir prim func.
🧰 Tools
🪛 Ruff (0.13.3)

11-11: Unused function argument: private

(ARG001)

🤖 Prompt for AI Agents
In tilelang/language/tir/entry.py around lines 1 to 40, the module docstring
still uses the old typing Union[PrimFunc, Callable]; update that docstring to
reflect the modern return type annotation by replacing any occurrence of
"Union[PrimFunc, Callable]" (or similar) with the new union syntax "PrimFunc |
Callable" so the docstring matches the function signature.

n_warmup: int = 1,
n_repeat: int = 1,
input_tensors: List[torch.Tensor] = None,
input_tensors: list[torch.Tensor] = None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Fix implicit Optional - missing | None in type annotation.

Line 226 has input_tensors: list[torch.Tensor] = None but should be input_tensors: list[torch.Tensor] | None = None. The function logic at lines 250 and 267 explicitly checks if input_tensors is None, so the type annotation must allow None.

Apply this diff:

-        input_tensors: list[torch.Tensor] = None,
+        input_tensors: list[torch.Tensor] | None = None,
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
input_tensors: list[torch.Tensor] = None,
input_tensors: list[torch.Tensor] | None = None,
🧰 Tools
🪛 Ruff (0.13.3)

226-226: PEP 484 prohibits implicit Optional

Convert to Optional[T]

(RUF013)

🤖 Prompt for AI Agents
In tilelang/profiler/__init__.py around line 226 the parameter annotation uses
`input_tensors: list[torch.Tensor] = None` which incorrectly omits the optional
union; update the annotation to `input_tensors: list[torch.Tensor] | None =
None` so the type system reflects that None is an allowed value (the function
later checks `if input_tensors is None` at lines ~250 and ~267).

LeiWang1999
LeiWang1999 previously approved these changes Oct 10, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants