[WY representation] Faster lower triangle inverse#289
Conversation
WalkthroughThis pull request introduces two new files that implement specialized Triton kernels: one for performing chunked scaled dot products and another for lower triangular matrix inversion. Additionally, new tests validate these functionalities, and an existing function in the delta rule module is refactored to utilize the new kernel functions with updated parameters. Changes
Sequence Diagram(s)sequenceDiagram
participant Caller as Caller
participant FPR as fwd_prepare_wy_repr
participant CSD as chunk_scaled_dot_kkt_fwd
participant ST as solve_tril
Caller->>FPR: Call fwd_prepare_wy_repr(k, v, beta, offsets, indices, head_first, chunk_size)
FPR->>CSD: Compute scaled dot product (A)
CSD-->>FPR: Return tensor A
FPR->>ST: Perform triangular inversion on A
ST-->>FPR: Return inverted matrix A
FPR-->>Caller: Return w, u, A
Possibly related PRs
Suggested reviewers
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (2)
🔇 Additional comments (12)
✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
tests/ops/test_solve_tril.py (2)
30-33: Clarify or re-check the skip condition for chunk-varlen tests.
The skip condition is triggered whenSKIP_TEST_CHUNK_VARLEN == "0", yet the reason states "TEST_CHUNK_VARLEN is enabled." This logic can be confusing. Consider ensuring the naming matches the intended behavior, such as skipping the test when the corresponding environment variable is actually enabled (truthy).
56-78: Reduce nested looping for improved test performance.
Intest_solve_tril_varlen, lines 72–77 implement a nested loop over sequence chunk boundaries. While correct, it may slow down large-scale test runs. Consider vectorized or batched approaches if performance becomes a bottleneck.fla/ops/utils/solve_tril.py (1)
25-71: Consider partial-chunk edge handling insolve_tril_16x16_kernel.
The kernel handles 16×16 blocks via a loop (line 63), assuming at least partial coverage. For smaller remaining chunks or minimal T, carefully verify that boundary conditions always hold. This can improve safety and clarity of partial-block logic.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
fla/ops/common/chunk_scaled_dot_kkt.py(1 hunks)fla/ops/utils/solve_tril.py(1 hunks)tests/ops/test_solve_tril.py(1 hunks)
🧰 Additional context used
🧬 Code Definitions (1)
tests/ops/test_solve_tril.py (2)
fla/ops/common/chunk_scaled_dot_kkt.py (1)
chunk_scaled_dot_kkt_fwd(69-99)fla/ops/utils/solve_tril.py (1)
solve_tril(234-281)
🪛 GitHub Actions: nvidia-h100-ci
tests/ops/test_solve_tril.py
[error] 43-43: TypeError: 'NoneType' object is not subscriptable
[error] 43-43: TypeError: 'NoneType' object is not subscriptable
[error] 43-43: TypeError: 'NoneType' object is not subscriptable
[error] 45-45: TypeError: 'NoneType' object is not subscriptable
[error] 45-45: TypeError: 'NoneType' object is not subscriptable
[error] 45-45: TypeError: 'NoneType' object is not subscriptable
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: test
🔇 Additional comments (5)
tests/ops/test_solve_tril.py (1)
34-53:Details
❓ Verification inconclusive
Investigate potential NoneType subscript issue in indexing.
The pipeline failures suggest a"TypeError: 'NoneType' object is not subscriptable"at lines 43 and 45. Possible causes include an unexpected zero or undefined value ofT, or an uninitializedAin certain environments.Use the following script to scan for scenarios where
Tmight be set to zero or incorrectly inferred:If any environment or code path sets T to 0, additional conditional guards or error handling might be required.
🏁 Script executed:
#!/bin/bash # Check for the presence or usage of T=0 in the codebase. rg "test_t_list" -A 3 rg -A3 "if compiled_mode"Length of output: 27139
Action: Confirm that
solve_trilalways returns a valid tensor before indexing its result.After reviewing the parameters across the codebase, we see that in this test the value for
Tis consistently set to nonzero values (e.g., 64, 128, 200, 300, 500). This indicates that the "NoneType" issue is unlikely to stem from an unexpected zero value ofTor an improperly computedAin this test. The reported error at lines 43 and 45 may instead be due to a code path inside thesolve_trilfunction returningNone.Please verify the following:
- Confirm that
solve_trildoes not returnNoneunder any circumstances.- Add input validation or an explicit guard in
solve_trilif there’s any possibility of a missing return value.🧰 Tools
🪛 GitHub Actions: nvidia-h100-ci
[error] 43-43: TypeError: 'NoneType' object is not subscriptable
[error] 43-43: TypeError: 'NoneType' object is not subscriptable
[error] 43-43: TypeError: 'NoneType' object is not subscriptable
[error] 45-45: TypeError: 'NoneType' object is not subscriptable
[error] 45-45: TypeError: 'NoneType' object is not subscriptable
[error] 45-45: TypeError: 'NoneType' object is not subscriptable
fla/ops/common/chunk_scaled_dot_kkt.py (2)
69-99: Looks solid for chunk-based scaled dot-product computations.
The high-level wrapper sets up shapes, preparesindices, and calls the Triton kernel. Implementation is clear and logically consistent with the design.
22-61:Details
❓ Verification inconclusive
Ensure robust handling of T, K, and chunk-size corner cases.
Insidechunk_scaled_dot_kkt_fwd_kernel, indexing usesT,K,BT,BKwith boundary checks. Ensure those checks protect against invalid or negative dimensions if unexpected values slip in, e.g., through input misconfiguration.You can confirm by searching for calls to
chunk_scaled_dot_kkt_fwdcombined with dimension checks:
🏁 Script executed:
#!/bin/bash rg -A 5 "chunk_scaled_dot_kkt_fwd("Length of output: 138
Action: Verify and enhance dimension validations in
chunk_scaled_dot_kkt_fwd_kernelThe review comment’s intent is still valid: ensure that the boundary checks for
T,K,BT, andBKare robust against unexpected or negative input dimensions. Note that the initial regex command produced a parse error because the parenthesis wasn’t escaped. To properly search for function calls in the codebase, please use an escaped pattern such as:rg -A 5 'chunk_scaled_dot_kkt_fwd\('Points to address:
- Confirm that all call sites passing values to
chunk_scaled_dot_kkt_fwd_kernelenforce proper validations of inputs (especially forT,K, and chunk sizes).- If these values can be misconfigured, consider adding explicit error handling or assertions at the entry point of the kernel.
- Re-run the corrected search command to verify that all calls to
chunk_scaled_dot_kkt_fwdutilize appropriate dimension checks.fla/ops/utils/solve_tril.py (2)
85-133: Merging 16×16 blocks into 32×32 blocks is well-structured.
The stepwise logic (lines 126–133) for merging blocks appears correct and well-bounded. The negative sign usage for inverting the sub-block accumulation aligns with the approach.
234-281: Implementation ofsolve_trilfunction meets stated objectives.
The function introduces chunked inverse logic for 16×16, 32×32, and 64×64 blocks, returning(I + A)^-1as promised. Assertions enforce matrix shape, data type, and contiguity. This is consistent with the documented contract.
zhiyuan1i
left a comment
zhiyuan1i
left a comment
There was a problem hiding this comment.
sonta, I think os.getenv("SKIP_TEST_CHUNK_VARLEN") needs to be 1 here, maybe I can make it more clear.
@pytest.mark.parametrize("H", test_h_list)
@pytest.mark.parametrize("cu_seqlens", test_t_varlen_list)
@pytest.mark.parametrize("chunk_size", [64, 32, 16])
@pytest.mark.skipif(
os.getenv("SKIP_TEST_CHUNK_VARLEN") == "1",
reason="Skipping test because TEST_CHUNK_VARLEN is enabled"
)
def test_solve_tril_varlen(H, cu_seqlens, chunk_size):
T = cu_seqlens[-1]
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
fla/ops/utils/solve_tril.py (2)
286-311: Test code looks good, but contains a commented-out breakpointThe test implementation correctly validates the implementation against a reference solution using PyTorch's inverse function.
However, there's a commented-out breakpoint on line 309 that should be removed before merging:
- # breakpoint()
235-239: Enhance docstring with parameter descriptionsThe docstring is clear about what the function does, but it would be helpful to add descriptions for each parameter, especially for
cu_seqlensandhead_firstwhich may not be self-explanatory to users unfamiliar with the codebase.def solve_tril(A, cu_seqlens=None, head_first=True, output_dtype=torch.float32): """ Compute the inverse of the lower triangular matrix A should be strictly lower triangular. Please make sure A.triu() == 0. return: (I + A)^-1 + + Parameters: + - A: Lower triangular matrix of shape [B, H, T, BT] if head_first=True else [B, T, H, BT] + - cu_seqlens: Optional. Cumulative sequence lengths for variable-length sequences + - head_first: If True, input is [B, H, T, BT], else [B, T, H, BT] + - output_dtype: Data type of the output matrix """
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
fla/ops/utils/solve_tril.py(1 hunks)tests/ops/test_solve_tril.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/ops/test_solve_tril.py
🧰 Additional context used
🪛 GitHub Actions: lint
fla/ops/utils/solve_tril.py
[error] 271-271: flake8: E231 missing whitespace after ','
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: test
🔇 Additional comments (3)
fla/ops/utils/solve_tril.py (3)
13-71: Matrix inversion kernel implementation looks goodThe
solve_tril_16x16_kernelfunction is well-implemented with proper autotuning configuration and boundary checks. The algorithm correctly computes the inverse of a 16x16 lower triangular matrix with appropriate memory layout handling.
73-133: Implementation of 32x32 inverse assembly kernel looks solidThe kernel effectively merges 16x16 inverse blocks into a 32x32 inverse matrix using the appropriate block matrix operations. Good use of boundary checks and type conversions.
135-232: 64x64 inverse assembly implementation is correctThe kernel correctly implements the block matrix operations required to merge 16x16 inverse blocks into a 64x64 inverse matrix. All boundary checks and type conversions are properly handled.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
fla/ops/utils/solve_tril.py (3)
135-145: Autotuning key includes unused parametersThe autotuning key includes parameters ('K', 'BK', 'BC') that don't appear in the function signature which may lead to suboptimal autotuning.
@triton.autotune( configs=[ triton.Config({}, num_warps=num_warps, num_stages=num_stages) for num_warps in [2, 4] for num_stages in [2, 3, 4] ], - key=['H', 'K', 'BT', 'BK', 'BC', 'HEAD_FIRST', 'USE_OFFSETS'], + key=['H', 'HEAD_FIRST', 'USE_OFFSETS'], )
234-243: Consider adding assertion to validate strictly lower triangular propertyThe function documentation states that
Ashould be strictly lower triangular, but there's no assertion to validate this property.def solve_tril(A, cu_seqlens=None, head_first=True, output_dtype=torch.float32): """ Compute the inverse of the lower triangular matrix A should be strictly lower triangular. Please make sure A.triu() == 0. return: (I + A)^-1 """ assert A.shape[-1] in [16, 32, 64] assert A.dtype == torch.float32, "A should be float32." assert A.is_contiguous(), "A should be contiguous." + assert torch.all(torch.triu(A, 1) == 0), "A should be strictly lower triangular (A.triu(1) == 0)." if head_first is True:
309-309: Remove debugging breakpointThere's a commented-out breakpoint statement which should be removed before the final submission.
from fla.ops.utils.testing import assert_close - # breakpoint() assert_close("solve_tril", A_inv, A_inv_ref, 0.0001)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
fla/ops/utils/solve_tril.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: test
🔇 Additional comments (5)
fla/ops/utils/solve_tril.py (5)
1-11: Clean and appropriate importsThe imports are appropriately organized and include all necessary libraries for implementing Triton kernels for matrix operations.
13-35: Well-configured kernel with proper autotuningThe
solve_tril_16x16_kernelhas appropriate Triton decorators:
- Effective heuristics for conditional execution
- Good autotuning configuration with varying warps and stages
- Clear parameter definitions with proper type annotations
36-71: Efficient implementation of 16x16 inverse calculationThe kernel is well-implemented with:
- Proper handling of program IDs for parallel execution
- Effective management of variable batch sizes and dimensions
- Correct use of Triton's block pointer functionality
- Appropriate boundary checking
- Numerically sound algorithm for inverting the lower triangular matrix
271-271: Corrected whitespace issueThe whitespace after the comma has been fixed, addressing the issue flagged in the previous review.
286-311: Comprehensive test caseThe test logic is thorough and provides good validation:
- Creates realistic test data
- Handles padding and reshaping appropriately
- Computes a reference solution using torch.inverse
- Compares results with suitable tolerance
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (4)
fla/ops/common/chunk_scaled_dot_kkt.py (4)
24-38: Consider adding more documentation or type hints.
While Triton constraints are often implicit, inline explanations of each parameter would improve clarity for future maintainers.
48-53: Boundary checks.
Only the first dimension is in boundary_check=(0,). Confirm if the second dimension also needs boundary enforcement when data is non-contiguous.
72-92: Chunk size tunability.
chunk_size=64is a practical default, but exploring other chunk sizes might yield better performance for different sequence lengths.
93-105: Automatic detection for USE_OFFSETS.
The heuristic-based approach is convenient but can be non-obvious to new contributors. Adding a short comment or docstring would help clarify.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
fla/ops/common/chunk_scaled_dot_kkt.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (2)
- GitHub Check: test
- GitHub Check: test
🔇 Additional comments (6)
fla/ops/common/chunk_scaled_dot_kkt.py (6)
1-3: File header looks good.
No issues spotted in the initial comments and license line.
4-11: Imports are appropriate.
No conflicts or missing imports identified for these lines.
13-23: Autotune and heuristics configuration.
The configs fornum_warpsandnum_stagesseem reasonable. Ensure performance is tested thoroughly across different data shapes.
39-47: Edge case handling for offsets.
Ifeos < bosever occurs, the effectiveTwould be negative, causing unexpected behavior. Consider an assertion or safe check.
54-63: Loop iteration bounds.
For very largeK, ensure the range infor i_k in range(tl.cdiv(K, BK))doesn't risk integer overflow or overly large loops.
64-69: Lower-triangular masking logic.
Masking out the upper triangle is correct for KKT or triangular-related computations.
yzhangcs
force-pushed
the
faster-deltanet
branch
from
389a0ad to
c2674c1
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
fla/ops/common/chunk_scaled_dot_kkt.py (2)
54-63: Potential performance consideration when multiplying large chunks.
Inside the loop, you computeb_kb = b_k * b_beta[:, None]and then dotl.dot(...). For very largeK, ensure that the partial sum fits in floating-point precision. If necessary, consider employing accumulation in a higher-precision type (e.g., FP32 or FP64) before storing the final result, to minimize numerical errors.
72-82: Docstring can elaborate on the lower-triangular constraint and offsets.
The docstring states "Compute beta * k * k^T," but the result is further processed to yield a strictly lower-triangular matrix. Consider mentioning this triangular masking step, and also how the optionalcu_seqlensoffsets play into chunked variable-length operations.- r""" - Compute beta * k * k^T - """ + r""" + Compute beta * k * k^T and then masks upper triangle to produce a strictly lower-triangular result. + If cu_seqlens is provided, the kernel uses chunk offsets for variable-length sequences. + """
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
fla/ops/common/chunk_scaled_dot_kkt.py(1 hunks)
🔇 Additional comments (5)
fla/ops/common/chunk_scaled_dot_kkt.py (5)
1-3: File header looks fine.
The UTF-8 declaration, copyright line, and spacing are properly set.
13-24: Triton decorators are well-organized.
The usage of@triton.heuristicsand@triton.autotunefor controlling runtime parameters looks correct. Specifying key configurations withnum_warpsandnum_stagesis consistent with Triton best practices to tune performance.
48-53: Verify correct pointer arithmetic forbetain HEAD_FIRST vs. non-HEAD_FIRST cases.
While the indexing logic forp_betaappears consistent, it could be helpful to document these pointer layouts more explicitly in the docstring or surrounding comments, to clarify why the stride changes between(T,)vs.(H,).Would you like a script to scan for all instances of
betausage across the repository to confirm consistent indexing?
64-69: Double-check the strictly lower-triangular masking.
Line 64 enforces all diagonal and upper-triangular elements to 0:b_A = tl.where(tl.arange(0, BT)[:, None] > tl.arange(0, BT)[None, :], b_A, 0)This zeroes out the diagonal as well. Confirm that this behavior (strictly lower-triangular) is intended, or whether the diagonal should be preserved.
92-105: Ensure all shape corner cases are covered.
The kernel launch uses:chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)](...)with
NT = triton.cdiv(T, BT)ifcu_seqlensisNone. Verify that extremely small or zero-sizedTdoes not cause off-by-one errors or boundary issues, particularly when the chunk size is greater thanT.
| i_t, i_bh = tl.program_id(0), tl.program_id(1) | ||
| i_b, i_h = i_bh // H, i_bh % H | ||
| if USE_OFFSETS: | ||
| i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32) | ||
| bos, eos = tl.load(offsets + i_n).to(tl.int32), tl.load(offsets + i_n + 1).to(tl.int32) | ||
| T = eos - bos | ||
| else: | ||
| bos, eos = i_b * T, i_b * T + T | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Consider renaming i_t to avoid overshadowing.
Here, i_t is read from tl.program_id(0), but then overwritten via the offsets logic:
i_n, i_t = tl.load(indices + i_t * 2), tl.load(indices + i_t * 2 + 1)This can create confusion and obscure the kernel’s indexing logic. A clearer naming convention for the chunk index (e.g., chunk_id) would improve maintainability.
Below is a sample diff showing how to rename the overwritten variable:
- i_t, i_bh = tl.program_id(0), tl.program_id(1)
- if USE_OFFSETS:
- i_n, i_t = tl.load(indices + i_t * 2).to(tl.int32), tl.load(indices + i_t * 2 + 1).to(tl.int32)
+ chunk_id, i_bh = tl.program_id(0), tl.program_id(1)
+ if USE_OFFSETS:
+ i_n, local_t = tl.load(indices + chunk_id * 2).to(tl.int32), tl.load(indices + chunk_id * 2 + 1).to(tl.int32)
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (5)
tests/ops/test_delta.py (1)
166-170: Ensure consistency in error threshold adjustments.
Similar to the earlier block, increasing the thresholds for these asserts can hide small, potentially important deviations. Consider verifying the root cause of any test failures that prompted this change.If necessary, refine kernels or intermediate calculations to improve numerical precision, rather than only loosening the test checks.
fla/ops/common/chunk_scaled_dot_kkt.py (2)
62-62: Assess potential mixed dtype in the dot product.
The code doesb_kb.to(b_k.dtype)prior totl.dot, which might cause unintended truncations. Consider verifying consistent dtypes or employing higher precision for intermediate sums if numeric stability is a priority.
104-127: Evaluate kernel tiling constants and performance.
The autotuning parameters (num_warps=[2,4,8],num_stages=[2,3,4],BK=64) look reasonable, but large K or T may still benefit from further kernel optimizations or dynamic block sizing.Do you want to explore a larger range of warp counts or advanced tile shapes for better GPU occupancy?
fla/ops/utils/solve_tril.py (2)
255-274: Consider adding explicit verification of the strictly lower triangular property.While the docstring states "A should be strictly lower triangular," there's no explicit verification of this requirement in the implementation, which could lead to incorrect results if a non-conforming matrix is provided.
Consider adding a verification step before processing:
def solve_tril( A: torch.Tensor, cu_seqlens: Optional[torch.Tensor] = None, head_first: bool = False, output_dtype: torch.dtype = torch.float ) -> torch.Tensor: """ Compute the inverse of the lower triangular matrix A should be strictly lower triangular, i.e., A.triu() == 0. ... """ assert A.shape[-1] in [16, 32, 64] assert A.dtype == torch.float, "A should be float32." + # Verify A is strictly lower triangular + if head_first: + is_tril = torch.all(torch.triu(A, diagonal=1) == 0) + else: + is_tril = torch.all(torch.triu(A, diagonal=1) == 0) + assert is_tril, "Input matrix must be strictly lower triangular (A.triu(diagonal=1) == 0)"
275-276: Consider relaxing dtype requirements for greater flexibility.Currently, the implementation strictly requires float32 input. Consider supporting other precision types (float16, bfloat16) that might be useful in different contexts, particularly for large models or memory-constrained environments.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
fla/ops/common/chunk_scaled_dot_kkt.py(1 hunks)fla/ops/delta_rule/wy_fast.py(2 hunks)fla/ops/utils/solve_tril.py(1 hunks)tests/ops/test_delta.py(3 hunks)
🧰 Additional context used
🧬 Code Definitions (1)
fla/ops/delta_rule/wy_fast.py (2)
fla/ops/common/chunk_scaled_dot_kkt.py (1)
chunk_scaled_dot_kkt_fwd(72-126)fla/ops/utils/solve_tril.py (1)
solve_tril(249-318)
🔇 Additional comments (14)
tests/ops/test_delta.py (2)
21-21: Acknowledge expanded test coverage with larger sequence length.
By including 512 in thetest_t_list, the tests now cover higher sequence lengths, which helps identify performance issues and edge cases. Ensure that necessary GPU memory is available and that test runtime remains acceptable for large T.
90-96: Relaxed error thresholds may mask subtle differences.
Raising the thresholds from 0.005/0.007 to 0.006/0.008 reduces the strictness of the test. Confirm the new thresholds are appropriate and won't overlook potential numerical instability in the kernel.Would you like to double-check numerical stability across different devices or random seeds to ensure these relaxed thresholds remain reliable?
fla/ops/delta_rule/wy_fast.py (4)
10-11: Good integration with chunked K-K^T and TRIL solver.
These new imports improve code organization by extracting specialized logic into separate modules.
223-230: Validate correct usage ofchunk_scaled_dot_kkt_fwd.
The function call is consistent with the docstring, but double-check that calling code supplies correct shapes and ensures data is contiguous if required.Do you want to run an additional script to confirm the shapes of
kandbetaacross the codebase?
231-236: Verify return dtype alignment insolve_tril.
ReturningAink.dtypeafter a TRIL solve might introduce subtle precision differences ifk.dtypeis lower thanfloat32.Confirm that
kis typically float32 or a higher-precision dtype. Alternatively, keep the returned matrix in float32 for stability.
220-220:Details
❓ Verification inconclusive
Changing
head_firstdefault to False.
Previously, the default was True. This can alter the shape assumptions in downstream code. Verify that all external usage correctly accounts for the new default.
🏁 Script executed:
#!/bin/bash # Search codebase for calls to `fwd_prepare_wy_repr` that rely on the old default. rg -A 3 "fwd_prepare_wy_repr\(.*\)"Length of output: 37
Verify Impact of Default Change on
head_firstThe default value for
head_firstinfwd_prepare_wy_reprhas been changed fromTruetoFalseinfla/ops/delta_rule/wy_fast.pyat line 220. Our automated search for calls tofwd_prepare_wy_reprthat might rely on the old default did not return any instances. However, this outcome is inconclusive. Please manually verify that all external usages either explicitly set thehead_firstparameter or operate correctly under the new default value.
- File to check:
fla/ops/delta_rule/wy_fast.py(line 220)- Action Required: Confirm that downstream code and external calls to
fwd_prepare_wy_reprcorrectly handle the change fromhead_first=Truetohead_first=False.fla/ops/common/chunk_scaled_dot_kkt.py (1)
39-47: Variable name overshadowing remains.
Reusingi_tto both denotetl.program_id(0)and the chunk offset may cause confusion.fla/ops/utils/solve_tril.py (7)
14-71: Well-optimized kernel implementation with clear structure.The kernel correctly implements an efficient algorithm for 16x16 lower triangular matrix inversion with proper boundary checks and memory access patterns. The autotuning configuration provides good performance optimization opportunities.
73-132: Effective 32x32 matrix inversion using block matrix approach.The implementation correctly uses block matrix inversion formulas to merge 16x16 inversions into a 32x32 inverse. The boundary checks and precision handling are properly implemented.
134-246: Comprehensive 64x64 matrix inversion with proper block structure.The implementation effectively handles the complexity of a 64x64 matrix inversion by decomposing it into appropriate subblocks. The matrix multiplication operations and boundary checks are properly sequenced to ensure correct results.
285-296: LGTM: Correct kernel selection and parameterization.The function correctly prepares chunk indices and calls the solve_tril_16x16_kernel with appropriate parameters for handling variable-length sequences.
304-317: LGTM: Proper kernel selection based on matrix dimensions.The code correctly selects the appropriate merge function based on the matrix dimensions (32x32 or 64x64) and passes all the necessary parameters.
285-286: Nice use of triton.cdiv for calculating grid dimensions.The calculation of
NTproperly handles both variable-length sequences (using indices length) and fixed-length sequences (using triton.cdiv), ensuring efficient kernel grid configuration.
248-318: Excellent orchestration function for different matrix sizes.The
solve_trilfunction provides a clean interface that encapsulates the complexity of the matrix inversion process. It correctly handles different matrix sizes (16x16, 32x32, 64x64) and memory layouts (head-first or not).
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
fla/ops/utils/solve_tril.py (1)
134-246: Comprehensive 64x64 matrix inversion with block decompositionThe implementation correctly extends the block inversion approach to 64x64 matrices, handling the more complex interdependencies between blocks. The use of IEEE precision and proper boundary checks is commendable.
Consider adding comments explaining the mathematical relationships between blocks to improve maintainability, as the block operations are quite dense and complex.
fla/ops/common/chunk_scaled_dot_kkt.py (1)
66-66: Consider optimizing out the upper-triangle computation.Currently, the entire
b_Ablock is computed, and then the upper triangular part is zeroed out. If only the lower triangle is of interest, compute it directly to improve efficiency or reduce unnecessary work.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
fla/ops/common/chunk_scaled_dot_kkt.py(1 hunks)fla/ops/utils/solve_tril.py(1 hunks)
🔇 Additional comments (7)
fla/ops/utils/solve_tril.py (5)
14-71: Well-implemented Triton kernel for 16x16 matrix inversionThe implementation of
solve_tril_16x16_kernelis robust, with proper boundary checks and memory layout handling. The kernel efficiently computes the inverse of a strictly lower triangular matrix, building the inverse iteratively and adding identity at the end.
73-132: Effective implementation of 32x32 matrix inversion via block mergingThis kernel correctly implements the block matrix inversion technique by combining 16x16 inverses into a 32x32 inverse. The use of
tl.dotwith explicitinput_precision='ieee'ensures numerical stability, which is crucial for matrix inversions.
248-318: Well-orchestrated function with clear interfaceThe
solve_trilfunction provides a clean interface to the underlying kernels with proper input validation, output type handling, and support for different layouts and variable sequence lengths.
286-286: Minor style issue with spacingThere's a missing space after the comma in the
triton.cdiv(T,16)call.- NT = len(indices) if cu_seqlens is not None else triton.cdiv(T,16) + NT = len(indices) if cu_seqlens is not None else triton.cdiv(T, 16)
275-276: Clear input validation with assertionsThe assertions effectively validate that the input matrix has a supported size (16, 32, or 64) and the correct data type (float32). This prevents unexpected behavior and provides clear error messages.
fla/ops/common/chunk_scaled_dot_kkt.py (2)
40-43: Avoid overshadowingi_t.Reassigning
i_tafter reading it fromtl.program_id(0)can cause confusion and reduces the clarity of the kernel’s indexing logic. Renaming the first usage (or the second) improves maintainability and aligns with prior review feedback.
63-64: Double-check precision when castingb_kb.Casting
b_betatob_k.dtypebefore the dot product might introduce floating-point precision loss ifb_kis a lower-precision format thanbeta. Ensure the desired numeric accuracy is maintained, potentially by upcasting both to a higher precision (e.g., float32) before the multiplication.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
fla/ops/utils/solve_tril.py (1)
277-278: Consider supporting multiple precision types beyond float32.The function currently restricts input to float32 only. Consider supporting other precision types (like float16) for better flexibility, especially in performance-critical applications where reduced precision might be acceptable.
-assert A.dtype == torch.float, "A should be float32." +assert A.dtype in [torch.float, torch.float16], f"A should be float32 or float16, got {A.dtype}."
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
fla/ops/utils/solve_tril.py(1 hunks)tests/ops/test_delta.py(5 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/ops/test_delta.py
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: test
- GitHub Check: test
- GitHub Check: test
🔇 Additional comments (7)
fla/ops/utils/solve_tril.py (7)
14-71: Well-implemented kernel for inverting 16x16 lower triangular matrices with thorough optimizations.The kernel implementation effectively uses Triton's autotuning capabilities to optimize performance across different workloads. The code properly handles boundary checking, memory layouts, and different data types. The mathematical approach for matrix inversion is sound.
73-133: Efficient implementation of 32x32 lower triangular matrix inversion using block-wise approach.This kernel effectively merges the results from 16x16 inversions to compute a 32x32 inverse. The implementation follows standard mathematical techniques for block matrix inversion, with proper boundary checking and memory handling.
135-248: Comprehensive implementation of 64x64 lower triangular matrix inversion with proper mathematical approach.The kernel handles the complex case of merging multiple 16x16 results into a 64x64 inverse. Despite the complexity, the implementation is mathematically sound and includes proper memory handling and boundary checks.
257-276: Clear and informative docstring.The docstring clearly explains the function's purpose, parameters, and return values. It provides necessary information for users to understand how to properly use the function.
288-288: Use consistent spacing around operators.The spacing in this line is inconsistent with the style used in the rest of the file.
-NT = len(indices) if cu_seqlens is not None else triton.cdiv(T, 16) +NT = len(indices) if cu_seqlens is not None else triton.cdiv(T, 16)
250-321: Well-designed orchestration function with good error handling.The
solve_trilfunction effectively orchestrates the appropriate kernels based on matrix size and handles different memory layouts and sequence lengths correctly. The error handling through assertions helps catch potential issues early.
1-322:Details
❓ Verification inconclusive
Verify numerical stability with test cases.
While the implementation is mathematically sound, it's crucial to verify numerical stability across different matrix sizes and conditions, especially for edge cases.
🏁 Script executed:
#!/bin/bash # Check if there are test cases for this new implementation # Look for test files related to solve_tril rg -l "solve_tril" --type py tests/Length of output: 64
Ensure Numerical Stability Through Comprehensive Test Coverage
The presence of
tests/ops/test_solve_tril.pyis a good sign. Please double-check that this file covers a diverse set of scenarios, including various matrix sizes (i.e. for BT values of 16, 32, and 64) and potential edge cases (e.g. minimal sizes or atypical input values) to confirm numerical stability across all conditions.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
.github/workflows/nvidia-h100.yml (2)
50-50: Verify the MAX_JOBS configuration for flash-attn buildThe command now installs
flash-attnwithMAX_JOBS=4and the--no-build-isolationflag. This setting appears intended to control parallel build jobs on CI runners. Please verify that4is optimal for the available CI resources or consider parameterizing this value if needed in different environments.
51-51: Simplify triton installation commandThe updated command
pip install -U tritonremoves the previously specified index URL and trusted host parameters. This simplification is acceptable provided that the default PyPI repository serves the appropriate Triton builds. It’s worthwhile to double-check that this change aligns with expected package sourcing requirements.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
.github/workflows/nvidia-4090.yml(1 hunks).github/workflows/nvidia-h100.yml(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- .github/workflows/nvidia-4090.yml
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: test
- GitHub Check: test
- GitHub Check: test
🔇 Additional comments (1)
.github/workflows/nvidia-h100.yml (1)
47-48: Streamlined pip installation commands for core packagesThe change to use the short
-Uflag (instead of--upgrade) for installing packages likepytest,setuptools,wheel,ninja, andtorchsimplifies the command while preserving the intended functionality and ensuring that the latest versions are installed. This improvement increases readability and maintains consistency with similar workflows.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tests/ops/test_dplr_delta.py (2)
317-320: Platform-specific test skip is well-implemented.This skip condition appropriately excludes the test on Intel platforms where Triton kernel execution is problematic. The reason message clearly explains why the test is being skipped.
Consider adding a link to an issue tracker or more details about the specific failure mode to help future maintainers understand the limitation better.
409-412: Consistent platform-specific test skipping.This skip condition correctly matches the one applied to the
test_chunkfunction, ensuring consistent treatment of Intel platforms across similar test cases.To reduce duplication, consider creating a reusable marker or helper for Intel platform skipping if more tests need similar treatment in the future:
intel_skip = pytest.mark.skipif( device_platform == 'intel', reason="Intel Triton Failure" ) # Then use it like: @intel_skip def test_chunk(...): ...
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
tests/ops/test_dplr_delta.py(3 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: test
- GitHub Check: test
- GitHub Check: test
🔇 Additional comments (1)
tests/ops/test_dplr_delta.py (1)
12-12: Import addition for platform-specific test handling.The import of
device_platformis correctly added to support the new platform-specific test skip conditions later in the file.
zhiyuan1i
force-pushed
the
faster-deltanet
branch
from
a8aa59e to
67d86a1
Compare
3 participants
Summary by CodeRabbit