Improve code readability and linting workflow #216
Conversation
…testing Improves linting workflow in Makefile Adds selective ruff fixups when not on main and keeps full style run for the default branch. Introduces dedicated targets for linting, formatting, and pytest to streamline the quality gate.
Rewraps long flash sparse attention helpers to follow style limits, clarifying imports, parameter mapping, and varlen handling without altering behavior
Allows mask helpers to accept None for window size and min dtype so callers can rely on dtype-dependent defaults instead of supplying explicit sentinel values
Refactors preprocessing and kernel helpers for clearer pointer math, explicit grid lambdas, and consistent naming (Indices) so the Triton autotuned kernels remain easier to reason about while keeping behavior unchanged.
Improves readability by reformatting long signatures, decorators, and tensor allocations, aligning with style guidelines while leaving behavior untouched.
github-actions
bot
assigned
Evanwu1125,
ftgreat,
SNHuan,
Thanksyy,
wubingheng111 and
zacliu2023
There was a problem hiding this comment.
Pull request overview
This PR enhances code readability and establishes a comprehensive linting workflow for the flash_sparse_attn project. The changes primarily focus on formatting improvements through automatic code formatting tools and introducing structured development workflows.
Key Changes:
- Added comprehensive Makefile targets for automated testing, linting, and formatting (ruff-based)
- Applied consistent code formatting across all Python files (line breaks, spacing, parentheses alignment)
- Improved API design by making certain parameters optional (e.g.,
min_dtype,window_sizeincreate_mask)
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| Makefile | Introduced new targets for code quality: test, style, quality, fixup, and modified_only_fixup for automated linting and testing |
| flash_sparse_attn/utils/padding.py | Applied formatting improvements with multi-line function calls and tuple assignments |
| flash_sparse_attn/utils/mask.py | Reformatted code and removed unused window_size parameter from relu_mask call; changed parameter types to Optional |
| flash_sparse_attn/integrations/modeling_flash_sparse_attention_utils.py | Extensive formatting improvements to enhance readability of complex function calls and conditionals |
| flash_sparse_attn/integrations/import_utils.py | Minor formatting and removed extra blank line |
| flash_sparse_attn/integrations/flash_sparse_attention.py | Reformatted multi-line comprehensions and function calls |
| flash_sparse_attn/flash_sparse_attn_triton.py | Applied consistent formatting to Triton kernel code with improved line breaks in complex expressions |
| flash_sparse_attn/flash_sparse_attn_interface.py | Reformatted function definitions and removed commented-out unused variables |
| flash_sparse_attn/flash_sparse_attn_flex.py | Minor formatting improvements for consistency |
| flash_sparse_attn/flash_dmattn_triton.py | Extensive formatting improvements; contains one critical bug (missing return statement in grid function) |
| flash_sparse_attn/init.py | Improved error message formatting and import organization |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| grid = lambda META: (triton.cdiv(seqlen_q, META["BLOCK_M"]), batch * nheads) | ||
|
|
||
| def grid(META): | ||
| ( |
There was a problem hiding this comment.
The grid function is missing a return statement. The tuple should be returned explicitly.
| ( | |
| return ( |
| window_size: Optional[int], | ||
| min_dtype: Optional[float], | ||
| block_size: Optional[int], | ||
| type: str = "topk", |
There was a problem hiding this comment.
The parameter name 'type' shadows the built-in Python function 'type'. Consider renaming to 'mask_type' or 'method' to avoid shadowing built-in names.
| attn_mask = attn_mask[:, :, :, : key.shape[-2]] | ||
| else: | ||
| attn_mask = torch.ones((batch, nheads, seqlen_q, seqlen_k), device=query.device, dtype=query.dtype) | ||
| attn_mask = torch.ones( |
There was a problem hiding this comment.
Variable attn_mask is not used.
| modified_py_files=$$(git diff --name-only main...HEAD | grep '\.py$$' || true); \ | ||
| if [ -n "$$modified_py_files" ]; then \ | ||
| echo "Checking/fixing files: $${modified_py_files}"; \ | ||
| $(PYTHON) -m ruff check $${modified_py_files} --fix; \ | ||
| $(PYTHON) -m ruff format $${modified_py_files}; \ |
There was a problem hiding this comment.
The modified_py_files variable is populated from git diff output and then expanded unquoted in the ruff invocations, which allows file names to inject additional shell tokens. An attacker who can introduce a file with a crafted name (e.g., containing ; rm -rf /) into a branch could cause arbitrary shell commands to execute on a developer’s machine when make modified_only_fixup is run. To fix this, ensure modified_py_files is safely passed as arguments (e.g., by proper quoting or using a safer mechanism that avoids raw shell token expansion of file names).
| modified_py_files=$$(git diff --name-only main...HEAD | grep '\.py$$' || true); \ | |
| if [ -n "$$modified_py_files" ]; then \ | |
| echo "Checking/fixing files: $${modified_py_files}"; \ | |
| $(PYTHON) -m ruff check $${modified_py_files} --fix; \ | |
| $(PYTHON) -m ruff format $${modified_py_files}; \ | |
| modified_py_files=$$(git diff --name-only -z main...HEAD -- '*.py' || true); \ | |
| if [ -n "$$modified_py_files" ]; then \ | |
| printf '%s' "$$modified_py_files" | while IFS= read -r -d '' file; do \ | |
| echo "Checking/fixing file: $$file"; \ | |
| $(PYTHON) -m ruff check "$$file" --fix; \ | |
| $(PYTHON) -m ruff format "$$file"; \ | |
| done; \ |
Summary
Root Cause
Changes
Reproduction
Tests
Compatibility
Checklist