[BC] Capitalize all envs by yzhangcs · Pull Request #650 · fla-org/flash-linear-attention

yzhangcs · 2025-11-19T06:55:37Z

Summary by CodeRabbit

Refactor
- Standardized internal configuration flag and constant naming conventions across modules to improve code consistency and maintainability.

coderabbitai · 2025-11-19T06:55:51Z

Caution

Review failed

The pull request is closed.

Walkthrough

This PR refactors feature detection flag constants across the codebase, renaming lowercase variable names (e.g., is_amd, is_nvidia_hopper, use_cuda_graph) to uppercase constant names (e.g., IS_AMD, IS_NVIDIA_HOPPER, USE_CUDA_GRAPH). It updates fla/utils.py to define and export uppercase constants, introduces lowercase aliases for backward compatibility, and replaces all import and usage references throughout the codebase.

Changes

Cohort / File(s)	Summary
Core utility module `fla/utils.py`	Renamed feature flag variables to uppercase constants: `supports_autotune_cache` → `SUPPORTS_AUTOTUNE_CACHE`, `is_amd` → `IS_AMD`, `is_intel` → `IS_INTEL`, `is_nvidia` → `IS_NVIDIA`, `is_intel_alchemist` → `IS_INTEL_ALCHEMIST`, `is_nvidia_hopper` → `IS_NVIDIA_HOPPER`, `use_cuda_graph` → `USE_CUDA_GRAPH`, `is_tf32_supported` → `IS_TF32_SUPPORTED`, `is_gather_supported` → `IS_GATHER_SUPPORTED`, `is_tma_supported` → `IS_TMA_SUPPORTED`. Added `_register_aliases()` function to create lowercase aliases pointing to new uppercase constants for backward compatibility.
Module activations and convolution kernels `fla/modules/activations.py`, `fla/modules/convolution.py`, `fla/modules/fused_bitlinear.py`, `fla/modules/grpo.py`, `fla/modules/l2norm.py`, `fla/modules/rotary.py`, `fla/modules/token_shift.py`	Updated imports to use `IS_AMD` instead of `is_amd`; updated `NUM_WARPS_AUTOTUNE` conditionals to reference the new uppercase constant.
Module KL divergence and cross-entropy `fla/modules/fused_kl_div.py`, `fla/modules/fused_linear_cross_entropy.py`	Updated imports to use `IS_NVIDIA_HOPPER` instead of `is_nvidia_hopper`; updated `STATIC_WARPS` conditionals to reference the new uppercase constant.
Common ops chunk kernels `fla/ops/common/chunk_delta_h.py`, `fla/ops/common/chunk_o.py`, `fla/ops/common/fused_chunk.py`	Updated imports to use `IS_NVIDIA_HOPPER` and `USE_CUDA_GRAPH` instead of `is_nvidia_hopper` and `use_cuda_graph`; updated `NUM_WARPS` conditionals and Triton autotune decorator parameters.
Delta rule and Wyfast ops `fla/ops/delta_rule/wy_fast.py`, `fla/ops/gated_delta_product/chunk_deltaproduct_h.py`, `fla/ops/gated_delta_product/chunk_deltaproduct_o.py`	Updated imports to use `IS_NVIDIA_HOPPER` and/or `USE_CUDA_GRAPH` instead of `is_nvidia_hopper` and `use_cuda_graph`; updated `NUM_WARPS` conditionals and autotune decorator parameters.
Generalized delta rule DPLR modules `fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py`, `fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py`, `fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py`, `fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py`, `fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py`, `fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py`	Updated imports to use `IS_AMD`, `IS_GATHER_SUPPORTED`, and `USE_CUDA_GRAPH` instead of `is_amd`, `is_gather_supported`, and `use_cuda_graph`; updated `NUM_WARPS_AUTOTUNE` conditionals, autotune decorators, and kernel parameter passing to use new uppercase constants.
Generalized delta rule DPLR fast variants `fla/ops/generalized_delta_rule/dplr/fused_recurrent.py`, `fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py`, `fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py`	Updated imports to use `IS_INTEL_ALCHEMIST`, `IS_GATHER_SUPPORTED`, and `USE_CUDA_GRAPH` instead of `is_intel_alchemist`, `is_gather_supported`, and `use_cuda_graph`; updated autotune configurations and kernel parameters accordingly.
Generalized delta rule IPLR modules `fla/ops/generalized_delta_rule/iplr/chunk.py`, `fla/ops/generalized_delta_rule/iplr/wy_fast.py`	Updated imports to use `IS_NVIDIA_HOPPER` and `USE_CUDA_GRAPH` instead of `is_nvidia_hopper` and `use_cuda_graph`; updated `NUM_WARPS` and autotune configurations.
KDA gate and wyfast ops `fla/ops/kda/gate.py`, `fla/ops/kda/wy_fast.py`	Updated imports to use `IS_AMD` and `IS_TF32_SUPPORTED` instead of `is_amd` and `is_tf32_supported`; updated `NUM_WARPS_AUTOTUNE` and `DOT_PRECISION` conditionals.
Mesa net chunk ops `fla/ops/mesa_net/chunk_h_kv_intra_bwd.py`, `fla/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py`	Updated imports to use `IS_NVIDIA_HOPPER` instead of `is_nvidia_hopper`; updated `NUM_WARPS` conditionals.
RWKV6 and RWKV7 ops `fla/ops/rwkv6/chunk.py`, `fla/ops/rwkv7/channel_mixing.py`, `fla/ops/rwkv7/fused_addcmul.py`, `fla/ops/rwkv7/fused_k_update.py`, `fla/ops/rwkv7/fused_recurrent.py`	Updated imports to use `IS_AMD` and `USE_CUDA_GRAPH` instead of `is_amd` and `use_cuda_graph`; updated `NUM_WARPS_AUTOTUNE` conditionals and autotune decorator parameters.
Simple GLA and TTT ops `fla/ops/simple_gla/parallel.py`, `fla/ops/ttt/fused_chunk.py`	Updated imports to use `IS_INTEL_ALCHEMIST` and `IS_NVIDIA_HOPPER` instead of `is_intel_alchemist` and `is_nvidia_hopper`; updated `NUM_WARPS` and `triton_config` conditionals.
Utility ops `fla/ops/utils/op.py`, `fla/ops/utils/softmax.py`, `fla/ops/utils/solve_tril.py`	Updated imports to use `IS_GATHER_SUPPORTED`, `IS_AMD`, and `IS_TMA_SUPPORTED` instead of `is_gather_supported`, `is_amd`, and `is_tma_supported`; updated conditional guard logic and kernel parameter passing.
Test files `tests/models/test_modeling_base.py`, `tests/modules/test_grpo.py`, `tests/modules/test_l2warp.py`, `tests/ops/test_comba.py`, `tests/ops/test_deltaformer.py`, `tests/ops/test_forgetting_attn.py`, `tests/ops/test_gated_delta.py`, `tests/ops/test_gated_delta_product.py`, `tests/ops/test_kda.py`, `tests/ops/test_mesa.py`, `tests/ops/test_path_attn.py`, `tests/ops/test_rwkv7.py`	Updated imports and skip conditions to use uppercase constants (`IS_AMD`, `IS_INTEL_ALCHEMIST`, `IS_NVIDIA_HOPPER`) instead of lowercase flag names.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Areas requiring extra attention:

fla/utils.py: The core aliasing mechanism (_register_aliases()) requires careful verification to ensure backward compatibility is properly maintained and that the alias creation doesn't introduce unexpected side effects.
fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py: This file adds GATHER_SUPPORTED as an explicit kernel parameter signature change, which is a structural modification beyond simple constant renaming and warrants closer inspection.
Complex module files (fla/modules/fused_bitlinear.py, fla/ops/common/chunk_delta_h.py, fla/ops/gated_delta_product/chunk_deltaproduct_*.py): Multiple flags are updated in these files; verify that all conditional branches and autotune decorator parameters are correctly updated.
Test coverage validation: Ensure that the aliasing fallback in fla/utils.py allows tests using the old lowercase names (if any remain) to continue working without failures.

Possibly related PRs

[TMA] Accelerate solve_tril with TMA descriptors[skip test] #550: Adds TMA support and makes_tensor_descriptor handling with overlapping changes to fla/ops/utils/solve_tril.py and TMA-related flag definitions that this PR renames (is_tma_supported → IS_TMA_SUPPORTED).
[Triton] use tl.gather to enhance performance #270: Modifies Triton gather/gather-support logic and kernel signatures (e.g., chunk_A_* with GATHER_SUPPORTED parameter) that align with this PR's updates to IS_GATHER_SUPPORTED flag usage.
Add KDA #621: Adds/modifies KDA ops modules (fla/ops/kda/gate.py, fla/ops/kda/wy_fast.py) with feature flags that this PR renames (is_amd → IS_AMD, is_tf32_supported → IS_TF32_SUPPORTED).

Poem

🐰 Constants now stand tall and proud,
No more whispers, loud and clear—
Uppercase flags shout to the cloud,
While aliases keep old friends near.
Refactored paths now brightly glow! ✨

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment
Commit unit tests in branch env

📜 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 f0eabcb and 1e78cb4.

📒 Files selected for processing (53)

fla/modules/activations.py (1 hunks)
fla/modules/convolution.py (1 hunks)
fla/modules/fused_bitlinear.py (2 hunks)
fla/modules/fused_kl_div.py (2 hunks)
fla/modules/fused_linear_cross_entropy.py (3 hunks)
fla/modules/grpo.py (2 hunks)
fla/modules/l2norm.py (1 hunks)
fla/modules/rotary.py (1 hunks)
fla/modules/token_shift.py (1 hunks)
fla/ops/common/chunk_delta_h.py (3 hunks)
fla/ops/common/chunk_o.py (1 hunks)
fla/ops/common/fused_chunk.py (1 hunks)
fla/ops/delta_rule/wy_fast.py (1 hunks)
fla/ops/gated_delta_product/chunk_deltaproduct_h.py (3 hunks)
fla/ops/gated_delta_product/chunk_deltaproduct_o.py (1 hunks)
fla/ops/generalized_delta_rule/dplr/chunk_A_bwd.py (4 hunks)
fla/ops/generalized_delta_rule/dplr/chunk_A_fwd.py (3 hunks)
fla/ops/generalized_delta_rule/dplr/chunk_h_bwd.py (2 hunks)
fla/ops/generalized_delta_rule/dplr/chunk_h_fwd.py (2 hunks)
fla/ops/generalized_delta_rule/dplr/chunk_o_bwd.py (4 hunks)
fla/ops/generalized_delta_rule/dplr/chunk_o_fwd.py (2 hunks)
fla/ops/generalized_delta_rule/dplr/fused_recurrent.py (2 hunks)
fla/ops/generalized_delta_rule/dplr/wy_fast_bwd.py (2 hunks)
fla/ops/generalized_delta_rule/dplr/wy_fast_fwd.py (5 hunks)
fla/ops/generalized_delta_rule/iplr/chunk.py (3 hunks)
fla/ops/generalized_delta_rule/iplr/wy_fast.py (1 hunks)
fla/ops/kda/gate.py (1 hunks)
fla/ops/kda/wy_fast.py (2 hunks)
fla/ops/mesa_net/chunk_h_kv_intra_bwd.py (1 hunks)
fla/ops/mesa_net/chunk_h_kv_intra_bwd_separate.py (1 hunks)
fla/ops/rwkv6/chunk.py (9 hunks)
fla/ops/rwkv7/channel_mixing.py (3 hunks)
fla/ops/rwkv7/fused_addcmul.py (6 hunks)
fla/ops/rwkv7/fused_k_update.py (1 hunks)
fla/ops/rwkv7/fused_recurrent.py (2 hunks)
fla/ops/simple_gla/parallel.py (1 hunks)
fla/ops/ttt/fused_chunk.py (1 hunks)
fla/ops/utils/op.py (2 hunks)
fla/ops/utils/softmax.py (1 hunks)
fla/ops/utils/solve_tril.py (2 hunks)
fla/utils.py (4 hunks)
tests/models/test_modeling_base.py (3 hunks)
tests/modules/test_grpo.py (2 hunks)
tests/modules/test_l2warp.py (3 hunks)
tests/ops/test_comba.py (4 hunks)
tests/ops/test_deltaformer.py (3 hunks)
tests/ops/test_forgetting_attn.py (2 hunks)
tests/ops/test_gated_delta.py (4 hunks)
tests/ops/test_gated_delta_product.py (4 hunks)
tests/ops/test_kda.py (5 hunks)
tests/ops/test_mesa.py (4 hunks)
tests/ops/test_path_attn.py (4 hunks)
tests/ops/test_rwkv7.py (3 hunks)

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.}

gemini-code-assist · 2025-11-19T06:56:25Z

Summary of Changes

Hello @yzhangcs, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request undertakes a significant refactoring effort to standardize the naming of global boolean flags related to hardware and environment capabilities within the fla library. By capitalizing these constants, the change improves code readability and adheres to common Python style guides for global constants. To mitigate the impact on existing codebases, a mechanism for backward compatibility has been implemented, providing seamless transition for users.

Highlights

Naming Convention Standardization: Global environment-related boolean flags (e.g., is_amd, is_nvidia_hopper, use_cuda_graph) have been renamed to follow a capitalized constant naming convention (e.g., IS_AMD, IS_NVIDIA_HOPPER, USE_CUDA_GRAPH) across numerous files in the fla library.
Backward Compatibility: A new internal _register_aliases function has been introduced in fla/utils.py to automatically create lowercase aliases for the newly capitalized constants, ensuring that existing code relying on the old naming convention continues to function without immediate breakage.
Copyright Headers: Several files have had their copyright headers added or updated to reflect the standard Copyright (c) 2023-2025, Songlin Yang, Yu Zhang.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request is a great refactoring effort that improves code style by capitalizing constants across the entire codebase. The changes are consistent and well-executed. I particularly appreciate the addition of the _register_aliases function in fla/utils.py to maintain backward compatibility for external users, which is a thoughtful touch. I've found one minor issue with a duplicated line of code, which I've commented on.

gemini-code-assist · 2025-11-19T07:03:38Z

        return oxr, oxw, oxk, oxv, oxa, oxg
    else:
        return oxr, oxw, oxk, oxv, oxa, None
+        return oxr, oxw, oxk, oxv, oxa, None


This return statement is a duplicate of the one on the previous line, making it unreachable code. It should be removed for code clarity and to avoid potential confusion.

[BC] Capitalize all envs

1e78cb4

yzhangcs merged commit 40dbae9 into main Nov 19, 2025
3 of 5 checks passed

yzhangcs deleted the env branch November 19, 2025 06:59

gemini-code-assist bot reviewed Nov 19, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[BC] Capitalize all envs#650

[BC] Capitalize all envs#650
yzhangcs merged 1 commit intomainfrom
env

yzhangcs commented Nov 19, 2025 •

edited by coderabbitai bot

Loading

Uh oh!

coderabbitai bot commented Nov 19, 2025 •

edited

Loading

Review failed

Uh oh!

gemini-code-assist bot commented Nov 19, 2025

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Nov 19, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

yzhangcs commented Nov 19, 2025 • edited by coderabbitai bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

coderabbitai bot commented Nov 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review failed

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Poem

Uh oh!

gemini-code-assist bot commented Nov 19, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Nov 19, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

yzhangcs commented Nov 19, 2025 •

edited by coderabbitai bot

Loading

coderabbitai bot commented Nov 19, 2025 •

edited

Loading