Repository navigation
[Refactor] Support auto swizzling for tma store and phaseout related layout annotations - #1509
Conversation
… sink examples and kernel analysis script for improved clarity and performance.
…to enhance code clarity and performance.
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
|
Warning Rate limit exceeded@LeiWang1999 has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 7 minutes and 0 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📒 Files selected for processing (2)
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. 📝 WalkthroughWalkthroughRemoved many explicit TileLang swizzled-layout annotations and several shared-allocation scope arguments across examples and benchmarks; introduced conservative 1D BulkLoad/BulkStore handling with 1D‑TMA guards and fallbacks in copy lowering; added enum-to-string debug helpers and replaced the two-arg linear-layout API with a shape-based linear layout. Changes
Sequence Diagram(s)(silently omitted) Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
examples/attention_sink/example_gqa_sink_bwd_bhsd.py (1)
11-17: LGTM: Simplified SM version handling.The change removes the exception for unsupported SM versions and uses a default configuration for SM >= 90, which aligns with the PR objectives. The larger config (block_M=128, num_stages=2, threads=256) is appropriate for newer architectures with more resources.
Consider adding a brief comment explaining why SM 80 uses different configuration parameters:
Optional: Add explanatory comment
def get_bwd_configs(): + # SM 80 (A100) uses smaller blocks due to resource constraints + # SM 90+ (H100, etc.) uses larger blocks for better performance sm_major, sm_minor = torch.cuda.get_device_capability() sm_version = sm_major * 10 + sm_minor if sm_version == 80: return 64, 32, 1, 128 else: return 128, 32, 2, 256examples/linear_attention/example_mamba_chunk_scan.py (1)
122-122: Consider removing redundant scope parameter for consistency.Line 122 still has an explicit
scope="shared.dyn"parameter, while similar allocations at lines 115, 117, 120, 123, and 128 had their scope parameters removed in this refactoring. Since"shared.dyn"is the default forT.alloc_shared, this explicit parameter is redundant.🔎 Proposed change for consistency
- x_shared = T.alloc_shared((block_K, block_N), dtype, scope="shared.dyn") + x_shared = T.alloc_shared((block_K, block_N), dtype)
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (59)
benchmark/mamba2/benchmark_mamba_chunk_scan.pybenchmark/matmul/benchmark_matmul.pybenchmark/matmul_fp8/benchmark_matmul.pyexamples/analyze/example_conv_analyze.pyexamples/attention_sink/example_gqa_sink_bwd_bhsd.pyexamples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.pyexamples/attention_sink/example_mha_sink_bwd_bhsd.pyexamples/attention_sink/example_mha_sink_fwd_bhsd.pyexamples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.pyexamples/bitnet-1.58b/kernel_benchmark/tilelang_bitnet_158_int8xint2_prefill.pyexamples/cast/example_group_per_split_token_cast_to_fp8.pyexamples/cast/example_per_token_cast_to_fp8.pyexamples/convolution/example_convolution.pyexamples/convolution/example_convolution_autotune.pyexamples/deepseek_mla/example_mla_decode.pyexamples/deepseek_mla/example_mla_decode_paged.pyexamples/deepseek_mla/example_mla_decode_persistent.pyexamples/deepseek_mla/experimental/example_mla_decode_kv_fp8.pyexamples/deepseek_nsa/benchmark/benchmark_nsa_fwd.pyexamples/deepseek_nsa/example_tilelang_nsa_bwd.pyexamples/deepseek_v32/sparse_mla_bwd.pyexamples/dequantize_gemm/example_dequant_gemm_bf16_fp4_hopper.pyexamples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper.pyexamples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper_tma.pyexamples/dequantize_gemm/example_dequant_gemm_fp4_hopper.pyexamples/dequantize_gemm/example_dequant_gemm_w4a8.pyexamples/dequantize_gemm/example_dequant_groupedgemm_bf16_mxfp4_hopper.pyexamples/dsa_sparse_finetune/sparse_mla_bwd.pyexamples/flash_attention/README.mdexamples/flash_attention/example_gqa_bwd.pyexamples/flash_attention/example_gqa_bwd_tma_reduce.pyexamples/flash_attention/example_gqa_bwd_tma_reduce_varlen.pyexamples/flash_attention/example_gqa_bwd_wgmma_pipelined.pyexamples/flash_attention/example_gqa_fwd_varlen.pyexamples/flash_attention/example_mha_bwd_bhsd.pyexamples/flash_attention/example_mha_bwd_bshd_wgmma_pipelined.pyexamples/flash_attention/example_mha_fwd_varlen.pyexamples/flash_decoding/example_gqa_decode.pyexamples/flash_decoding/example_gqa_decode_varlen_logits.pyexamples/flash_decoding/example_mha_inference.pyexamples/gdn/example_chunk_delta_bwd.pyexamples/gdn/example_chunk_delta_h.pyexamples/gdn/example_chunk_o.pyexamples/gdn/example_chunk_o_bwd.pyexamples/gdn/example_chunk_scaled_dot_kkt.pyexamples/gdn/example_wy_fast.pyexamples/linear_attention/example_linear_attn_bwd.pyexamples/linear_attention/example_linear_attn_fwd.pyexamples/linear_attention/example_mamba_chunk_scan.pyexamples/linear_attention/example_mamba_chunk_state.pyexamples/minference/example_vertical_slash_sparse_attn.pyexamples/warp_specialize/example_warp_specialize_flashmla.pysrc/op/copy.ccsrc/op/copy.hsrc/op/gemm.hsrc/op/operator.htesting/python/jit/test_tilelang_jit_nvrtc.pytesting/python/jit/test_tilelang_jit_tvm_ffi.pytesting/python/language/test_tilelang_language_vectorized_cast.py
💤 Files with no reviewable changes (45)
- examples/dequantize_gemm/example_dequant_gemm_w4a8.py
- examples/minference/example_vertical_slash_sparse_attn.py
- examples/gdn/example_chunk_delta_h.py
- examples/deepseek_mla/example_mla_decode_persistent.py
- examples/gdn/example_chunk_o_bwd.py
- examples/flash_decoding/example_gqa_decode_varlen_logits.py
- examples/convolution/example_convolution.py
- examples/flash_decoding/example_gqa_decode.py
- examples/warp_specialize/example_warp_specialize_flashmla.py
- testing/python/jit/test_tilelang_jit_nvrtc.py
- examples/linear_attention/example_linear_attn_fwd.py
- examples/dequantize_gemm/example_dequant_gemm_bf16_fp4_hopper.py
- examples/dsa_sparse_finetune/sparse_mla_bwd.py
- examples/deepseek_nsa/benchmark/benchmark_nsa_fwd.py
- examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper.py
- benchmark/matmul_fp8/benchmark_matmul.py
- examples/dequantize_gemm/example_dequant_gemm_bf16_mxfp4_hopper_tma.py
- examples/flash_attention/example_mha_bwd_bshd_wgmma_pipelined.py
- examples/flash_attention/example_gqa_fwd_varlen.py
- examples/gdn/example_chunk_o.py
- examples/deepseek_v32/sparse_mla_bwd.py
- examples/convolution/example_convolution_autotune.py
- examples/dequantize_gemm/example_dequant_groupedgemm_bf16_mxfp4_hopper.py
- examples/deepseek_mla/example_mla_decode.py
- examples/flash_attention/example_mha_fwd_varlen.py
- testing/python/jit/test_tilelang_jit_tvm_ffi.py
- examples/attention_sink/example_mha_sink_fwd_bhsd_wgmma_pipelined.py
- examples/flash_attention/example_mha_bwd_bhsd.py
- examples/cast/example_per_token_cast_to_fp8.py
- examples/deepseek_mla/example_mla_decode_paged.py
- examples/attention_sink/example_mha_sink_fwd_bhsd.py
- examples/attention_sink/example_mha_sink_bwd_bhsd.py
- examples/flash_attention/example_gqa_bwd.py
- examples/flash_attention/example_gqa_bwd_tma_reduce.py
- examples/gdn/example_wy_fast.py
- examples/dequantize_gemm/example_dequant_gemm_fp4_hopper.py
- examples/linear_attention/example_linear_attn_bwd.py
- examples/gdn/example_chunk_scaled_dot_kkt.py
- examples/flash_attention/example_gqa_bwd_wgmma_pipelined.py
- examples/deepseek_nsa/example_tilelang_nsa_bwd.py
- examples/flash_attention/example_gqa_bwd_tma_reduce_varlen.py
- examples/deepseek_mla/experimental/example_mla_decode_kv_fp8.py
- examples/analyze/example_conv_analyze.py
- examples/gdn/example_chunk_delta_bwd.py
- benchmark/matmul/benchmark_matmul.py
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
Learnt from: botbw
Repo: tile-ai/tilelang PR: 691
File: src/tl_templates/cuda/gemm_sp_sm80.h:81-85
Timestamp: 2025-09-15T10:51:06.985Z
Learning: In CUTLASS tensor operation layouts, crosswise constants should be computed using sizeof(T) (bytes), not cutlass::sizeof_bits<T>::value (bits). This is the established pattern in the official CUTLASS codebase, as seen in default_mma_core_sparse_sm80.h.
Learnt from: botbw
Repo: tile-ai/tilelang PR: 691
File: src/tl_templates/cuda/gemm_sp_sm80.h:81-85
Timestamp: 2025-09-15T10:51:06.985Z
Learning: In CUTLASS tensor operation layouts, crosswise constants should be computed using sizeof(T) (bytes), not cutlass::sizeof_bits<T>::value (bits). However, the layout template parameter should use sizeof_bits<T>::value (bits). This is the established pattern in the official CUTLASS codebase, as seen in default_mma_core_sparse_sm80.h where Crosswise uses sizeof(ElementA) but the layout template uses sizeof_bits<ElementA>::value.
📚 Learning: 2025-11-14T07:56:11.098Z
Learnt from: lucifer1004
Repo: tile-ai/tilelang PR: 1256
File: testing/python/jit/test_tilelang_jit_gemm_nvrtc.py:55-115
Timestamp: 2025-11-14T07:56:11.098Z
Learning: In `testing/python/jit/test_tilelang_jit_gemm_nvrtc.py`, the global function `tilelang_callback_cuda_postproc` registered via `tvm.register_global_func(..., override=True)` is intentionally not restored after the test completes, as the persistent behavior is expected.
Applied to files:
testing/python/language/test_tilelang_language_vectorized_cast.py
🧬 Code graph analysis (7)
testing/python/language/test_tilelang_language_vectorized_cast.py (1)
tilelang/language/v2/dtypes.py (2)
float16(294-294)float4_e2m1fn(385-385)
examples/linear_attention/example_mamba_chunk_state.py (2)
tilelang/language/annotations.py (1)
annotate_layout(26-37)tilelang/layout/swizzle.py (1)
make_swizzled_layout(62-71)
src/op/gemm.h (1)
tilelang/tileop/gemm/__init__.py (1)
GemmInst(33-52)
examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py (1)
tilelang/env.py (1)
disable_cache(272-273)
benchmark/mamba2/benchmark_mamba_chunk_scan.py (1)
tilelang/language/allocate.py (2)
alloc_shared(40-55)alloc_fragment(72-85)
examples/cast/example_group_per_split_token_cast_to_fp8.py (1)
tilelang/language/v2/dtypes.py (1)
int32(240-240)
examples/linear_attention/example_mamba_chunk_scan.py (1)
tilelang/language/allocate.py (2)
alloc_shared(40-55)alloc_fragment(72-85)
⏰ 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). (3)
- GitHub Check: Test for Python 3.12 with Nightly-ROCm-7.1 (on self-hosted-amd)
- GitHub Check: Test for Python 3.12 with Metal (on macos-latest)
- GitHub Check: Test for Python 3.12 with CUDA-12.8 (on self-hosted-nvidia)
🔇 Additional comments (15)
examples/flash_attention/README.md (1)
1-109: Documentation update aligns well with the refactoring objectives.The removal of explicit
T.annotate_layoutcalls for shared buffers (Q_shared, K_shared, V_shared) is consistent with the PR's goal of simplifying layout management by supporting auto-swizzling. The documented example remains accurate and functional—the buffers are still properly allocated in shared memory, just without manual layout annotations. The code flow and kernel logic remain unchanged and well-explained.examples/cast/example_group_per_split_token_cast_to_fp8.py (3)
34-34: LGTM: Cleaner scalar allocation.The change from
T.alloc_fragment((1,), T.int32)toT.alloc_var(dtype=T.int32)appropriately simplifies the allocation for a scalar offset value, improving code readability.
36-38: LGTM: Consistent scalar usage.The initialization and accumulation logic correctly use the scalar variable. The loop properly computes the cumulative offset by summing batch sizes from index 0 to
bg-1.
41-41: LGTM: Simplified array indexing.The slice expression correctly uses the scalar offset, removing the unnecessary
[0]indexing and improving readability.examples/linear_attention/example_mamba_chunk_scan.py (1)
115-115: LGTM - Scope parameter removal aligns with PR objectives.The removal of explicit
scopeparameters from theseT.alloc_sharedcalls is correct. They now default toscope="shared.dyn"as defined in thealloc_sharedfunction signature, which is the intended behavior for this refactoring.Also applies to: 117-117, 120-120, 123-123, 128-128
testing/python/language/test_tilelang_language_vectorized_cast.py (1)
130-131: Revert to standard test harness pattern.The test harness
tilelang.testing.main()is commented out and replaced with a hard-coded call to a single test function. This deviates from the standard pattern used by all other test files in the repository, which calltilelang.testing.main()to discover and run tests via pytest.Restore line 130 to
tilelang.testing.main()and remove the hard-coded test call on line 131.Likely an incorrect or invalid review comment.
src/op/operator.h (1)
35-47: LGTM!Clean utility function for debugging. All current
InferLevelenum cases are covered, and the default case handles potential future enum additions gracefully.src/op/copy.h (1)
31-55: LGTM!All
CopyInstenum values are covered with appropriate string representations. Consistent pattern with other*ToStringhelpers added in this PR.src/op/gemm.h (1)
25-58: LGTM!Both helper functions correctly cover all enum cases for
GemmWarpPolicyTypeandGemmInst. The string representations align with the Python enum's__repr__behavior.src/op/copy.cc (1)
548-573: LGTM on 1D TMA condition handling.The expanded condition correctly handles the new
kBulkLoad1DandkBulkStore1Dvariants. The early return for single-dimension 1D TMA store is appropriate since swizzling cannot be applied to truly 1D layouts.examples/flash_decoding/example_mha_inference.py (1)
183-184: LGTM!Simplifying the data path by copying directly from
Output_partialtopo_localeliminates the intermediate shared buffer (po_shared). This aligns with the PR's objective of removing explicit layout annotations and relying on auto-inferred layouts.examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py (1)
146-149: The simplified shared buffer allocations are correct and align with the PR objective of supporting automatic layout inference for TMA operations. The compiler's layout inference assigns appropriate layouts to each buffer based on context, and swizzling strategies are automatically handled by TileLang, resulting in no visible differences in the code. This approach reduces verbosity while maintaining performance through compiler-driven optimization.benchmark/mamba2/benchmark_mamba_chunk_scan.py (2)
251-256: Clarify whether layout annotations should be fully removed.The PR objectives indicate removing explicit swizzled layouts and layout annotations across the codebase. However, this file still retains
T.annotate_layoutwithmake_swizzled_layoutcalls forcb_sharedandx_residual_shared.Is this intentional? Should these annotations also be removed as part of the auto swizzling support, or do these specific tensors require explicit layout specifications?
228-241: Code state is correct:T.alloc_sharedcalls properly default to dynamic shared memory.All
T.alloc_sharedcalls in lines 228-241 correctly omit explicit scope parameters, relying on the defaultscope="shared.dyn"which supports auto swizzling. Since the kernel uses float16 and float32 dtypes (not bool), the dynamic scope applies uniformly across all shared memory allocations.Layout annotations (lines 251-256) are appropriately applied only to
cb_sharedandx_residual_shared, consistent with the swizzling optimization strategy.examples/linear_attention/example_mamba_chunk_state.py (1)
92-92: Clarify the rationale for asymmetric layout annotation removal.Only
x_sharedretains its swizzled layout annotation (line 92), whileacc_o_sharedhas no annotation. Sincex_sharedis used in load operations (lines 104, 112) andacc_o_sharedin store operations (lines 125-129), please confirm:
- Is this asymmetry intentional—does auto-swizzling now cover store paths to
acc_o_sharedwhile loads still require manual annotation forx_shared?- Has removing the
acc_o_sharedannotation been performance-validated against the previous version?- Are both buffers covered by the auto-optimization mentioned in the PR title, or is there a plan to eventually remove the
x_sharedannotation once auto-swizzling fully covers load operations?
| tilelang.disable_cache() | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check for cache disable patterns across attention examples
# Search for disable_cache calls in example files
echo "=== Searching for disable_cache in examples ==="
rg -n "disable_cache" examples/
echo -e "\n=== Searching for cache-related calls in attention examples ==="
fd "example.*attention.*\.py$" examples/ 2>/dev/null | head -20
echo -e "\n=== Check for cache-related patterns in all example files ==="
rg -l "cache" examples/ --type py | head -10
echo -e "\n=== Check context around the file in question ==="
if [ -f "examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py" ]; then
head -20 examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py
fiRepository: tile-ai/tilelang
Length of output: 2264
🏁 Script executed:
#!/bin/bash
# Check full file structure and autotune usage
echo "=== Full file content ==="
cat examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py
echo -e "\n=== Check if test imports this file ==="
grep -r "example_gqa_sink_fwd_bhsd_wgmma_pipelined" examples/ tests/
echo -e "\n=== Check if this file has __main__ guard ==="
grep -n "__main__" examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py
echo -e "\n=== Compare with flash_attention example (which also has disable_cache) ==="
head -25 examples/flash_attention/example_gqa_bwd_tma_reduce.pyRepository: tile-ai/tilelang
Length of output: 16540
Move cache disable to controlled execution scope.
The cache is disabled at module import time, creating a global side effect for all code that imports this module—including test files. Since the file has a __name__ == "__main__" guard, move tilelang.disable_cache() into that block or into the benchmark/tune functions that require it, rather than executing it unconditionally on import.
🤖 Prompt for AI Agents
In examples/attention_sink/example_gqa_sink_fwd_bhsd_wgmma_pipelined.py around
lines 13-14, tilelang.disable_cache() is called at import time causing a global
side effect; move this call into a controlled execution scope such as inside the
existing if __name__ == "__main__" block or into the specific benchmark/tune
functions that need it so importing the module does not change global cache
state. Ensure the call executes only when running the script or when the
specific functions are invoked, not during module import.
|
@regression-perf |
Performance Benchmark ReportTriggered by: @LeiWang1999 Results
Artifacts
|
This change introduces a layout annotation for the B_shared tensor in the example_dequant_gemm_bf16_mxfp4_hopper.py file, enhancing the memory layout optimization for better performance during matrix multiplication.
…larity and performance.
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…dimension - Renamed `makeGemmLayoutLinear` to `makeLinearLayout` and updated its implementation to handle arbitrary dimensions. - Updated related function calls in `gemm_layouts.cc` and `layout.cc` to use the new layout function. - Enhanced layout inference in `CumSumOpNode` to enforce linear layout for shared buffers in strict mode.
… arbitrary dimensions - Updated the function signature to take a `Buffer`, `BufferLoad`, or `BufferRegion` directly. - Simplified the implementation by removing argument checks and directly obtaining the shape from the buffer info. - Enhanced the documentation to clarify the function's purpose and usage.
|
@regression-perf |
Performance Benchmark ReportTriggered by: @LeiWang1999 Results
Artifacts
|
…g unused code related to dequantization GEMM example.
…/tilelang into refactor_layout_1223
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| import tilelang | ||
| import tilelang.language as T | ||
| from tilelang import tvm as tvm | ||
| from tvm import DataType | ||
| from tvm import tir |
There was a problem hiding this comment.
Deleting this module means examples/dequantize_gemm/test_example_dequantize_gemm.py and examples/dequantize_gemm/regression_example_dequantize_gemm.py still import example_dequant_gemm_bf16_mxfp4_hopper_tma, so running those tests/regressions now raises ModuleNotFoundError during import (before any CUDA checks). If the example is intentionally removed, the corresponding imports and test functions should be removed/updated; otherwise keep a stub or restore the module.
Useful? React with 👍 / 👎.
This pull request removes the usage of swizzled memory layouts and related annotations throughout the codebase, simplifying tensor allocation and layout management. It also makes minor improvements to variable allocation and thread binding in some kernels. These changes help streamline the code and reduce unnecessary complexity around layout specification.
Memory layout and annotation simplification:
Removed all calls to
T.annotate_layoutand the use ofmake_swizzled_layoutfor shared tensors in convolution, attention, and casting examples, as well as benchmark scripts. This affects files such asexample_conv_analyze.py,example_convolution.py,example_convolution_autotune.py,example_gqa_sink_bwd_bhsd.py,example_gqa_sink_fwd_bhsd_wgmma_pipelined.py,example_mha_sink_bwd_bhsd.py,example_mha_sink_fwd_bhsd.py,example_mha_sink_fwd_bhsd_wgmma_pipelined.py,example_deepseek_mla/example_mla_decode.py,example_deepseek_mla/example_mla_decode_paged.py, andexample_deepseek_mla/example_mla_decode_persistent.py. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] [12] [13] [14] [15] [16] [17] [18] [19] [20]Removed imports of
make_swizzled_layoutandtilelang.layoutwhere they are no longer needed, further cleaning up the code. [1] [2] [3] [4] [5] [6]Tensor allocation and scope changes:
benchmark_mamba_chunk_scan.pyto remove the explicitscope="shared"andscope="shared.dyn"parameters, using simpler allocation calls instead.Kernel and utility improvements:
Changed thread binding in
tilelang_bitnet_158_int8xint2_prefill.pyto useT.get_thread_binding(0)instead ofT.thread_binding(0, threads, "threadIdx.x"), simplifying thread management.Improved variable allocation in
example_group_per_split_token_cast_to_fp8.pyby replacing fragment allocation forrow_offsetwith a simple scalar variable, and updated its usage accordingly.Miscellaneous:
get_bwd_configs()inexample_gqa_sink_bwd_bhsd.pyto always return the default configuration for SM version 90 and above, removing the exception for unsupported versions.Summary by CodeRabbit
Refactor
Bug Fixes
Chores
Tests
✏️ Tip: You can customize this high-level summary in your review settings.