mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-24 07:04:00 +08:00
6873e19da2b267d3cb7d0305561ba0bed63feae3
45
Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 Lei WangandClaude Opus 4.6
|
3ee0988c86 |
[Refactor][Pipeline] Run pipeline rewriting before layout inference and stabilize tiled WS (#2002)
* Refactor optimization pipeline in phase.py to streamline barrier handling
- Consolidated the handling of shared barriers and pipeline planning by removing redundant conditional checks.
- Ensured that `LowerSharedBarrier`, `PipelinePlanning`, and `InjectSoftwarePipeline` are consistently applied, enhancing the clarity and efficiency of the optimization process.
This change improves the maintainability of the code while preserving existing functionality.
* Enhance tile operation handling with instruction annotation and refactor gemm.h
- Added a new `InstructionAnnotation` pass to annotate tile operations with their instruction kind before layout inference, improving the optimization pipeline's ability to reason about instruction mixes.
- Refactored `gemm.h` to move the `allowTcgen5Mma` and `allowWgmma` methods under the private section, enhancing code organization and encapsulation.
These changes improve the clarity and maintainability of the code while preserving existing functionality.
* Enhance multi-version buffer handling and introduce tile-level warp specialization
- Updated `multi_version_buffer_rewriter.cc` to improve read/write access detection for tile operations by analyzing `tl.tileop.region` calls, ensuring accurate buffer access tracking.
- Modified `phase.py` to integrate `ProducerConsumerWarpSpecializedTiled` before layout inference, allowing for high-level tile-op IR transformations that enhance producer/consumer splits.
- Added a new `ProducerConsumerWarpSpecializedTiled` function in `__init__.py` to facilitate tile-level warp specialization, improving the optimization pipeline's efficiency.
These changes enhance the handling of multi-version buffers and optimize the transformation process for tiled operations.
* Enable tiled WS for stage-1 pipelines
* Move Hopper pipeline planning before layout inference
* Run pipeline rewriting before layout inference
* Add pipeline refactor WIP: subtree_modified_ guard, debug prints, and plan
- inject_pipeline.cc: guard reads/writes recalculation with subtree_modified_
flag to prevent local.var buffer promotion to kernel parameters
- phase.py: add temporary debug prints after PipelinePlanning and
InjectSoftwarePipeline (to be removed during CI fix work)
- example_group_per_split_token_cast_to_fp8.py: add disable_cache for debugging
- docs/plan.md: implementation plan for fixing CI test failures
- draft.md: original design draft
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Remove debug artifacts: prints in phase.py, disable_cache in example
- Remove print("After PipelinePlanning"), print(mod),
print("After InjectSoftwarePipeline"), print(mod) from
LowerAndLegalize in tilelang/engine/phase.py
- Remove tilelang.disable_cache() from
examples/cast/example_group_per_split_token_cast_to_fp8.py
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix LayoutInference crash on multi-versioned pipeline buffers
After InjectSoftwarePipeline, multi-versioned buffers share the same
data Var as the original but have an extra leading dimension (num_stages).
LayoutInference's alias propagation and annotation handling tried to
Reshape layouts between these buffers, which failed because the total
element counts differ.
Guard three Reshape call sites in layout_inference.cc to skip sibling
buffers whose total storage size is incompatible with the source layout.
This lets multi-versioned buffers get their own layout inference instead
of inheriting an incompatible layout from the original buffer.
Fixes compilation failures in dequantize_gemm, GDN, and other kernels
that use software pipelining with shared memory buffers.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Strip pipeline annotations on WS fallback for TMA kernels
When ProducerConsumerWarpSpecializedTiled identifies a TMA kernel as a
warp-specialization candidate but the tiled rewriter cannot handle it
(e.g., conditional loop bodies like sparse block masks), the fallback
previously returned the original function with num_stages annotations
intact. PipelinePlanning and InjectSoftwarePipeline would then generate
non-WS TMA pipeline code with broken barrier phase tracking for
conditional pipeline bodies (barrier waits outside conditionals cause
deadlocks when the condition is false).
Fix: on WS fallback, strip num_stages annotations from pipeline loops
so that the pipeline passes skip the function. The kernel runs
unpipelined but correctly.
Fixes CUDA_ERROR_LAUNCH_FAILED in blocksparse_gemm and related TMA
kernels with conditional loop bodies.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add guarded phase-counter support to tiled WS pass
Port PhaseCounter and StageExprReplacer from the legacy
ProducerConsumerWarpSpecialized pass into the tiled WS pass to handle
conditional loop bodies (e.g., sparse block masks).
When the pipeline loop body is wrapped in an IfThenElse without else:
1. Unwrap the condition before classifying statements
2. Create separate producer/consumer PhaseCounters (local int32 buffers)
3. Use counter-based stage/parity expressions instead of loop-variable
4. Wrap producer and consumer bodies in the original condition
5. Increment counters at end of each guarded iteration
6. Rewrite shared-buffer stage indices via StageExprReplacer
This ensures barrier parity stays correct when iterations are
conditionally skipped, fixing CUDA_ERROR_LAUNCH_FAILED in
blocksparse_gemm and related TMA kernels with conditional execution.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add GemmSPNode support to pipeline planning and injection
Add GemmSPNode handling to:
- inject_pipeline.cc AddReadsWritesForTileOp: model A, E, B as reads
and C as write (E is the sparse metadata buffer)
- pipeline_planning.cc: same access model for dependency analysis
This makes sparse GEMM visible to the pipeline machinery for correct
stage assignment and buffer multi-versioning. However, the tile-op's
consumer-side buffer accesses still don't get stage-indexed because
the pipeline body rewriter can't rewrite high-level tile-op Call
arguments (an architectural limitation of running InjectSoftwarePipeline
before LowerTileOp).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix GemmSPNode::Lower to use MakeAccessPtrFromRegion
Replace whole-buffer access_ptr(1) calls with MakeAccessPtrFromRegion
for A, B, C, and E buffers in sparse GEMM lowering. This preserves
stage-specific region offsets from pipeline multi-versioning, matching
the dense GemmNode::Lower pattern.
CUDA output now shows correct stage-indexed consumer accesses:
gemm_sp_ss(..., (k%3)*8192, (k%3)*8192+27648, C_local, (k%3)*2048+49152)
instead of always using stage-0 offsets.
Note: gemm_sp still produces incorrect results because the kernel needs
warp specialization but TiledWSCandidate::Check doesn't recognize it
as a TMA candidate. The non-WS pipeline path generates structurally
different code from the reference WS path.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix gemm_sp and seer_attention regressions
Three changes in producer_consumer_ws_tiled.cc:
1. Require num_stages >= 2 for WS candidacy (was >= 1). Single-stage
kernels like seer_attention don't need WS and the transformation
produces incorrect results for them.
2. Add HasTmaPipeline() check to detect TMA kernels with pipeline
annotations that are rejected by the full WS candidate check (e.g.,
kernels with manual layout annotations like gemm_sp).
3. Strip num_stages annotations for rejected TMA pipeline kernels to
prevent InjectSoftwarePipeline from generating broken non-WS TMA
pipeline code.
One change in gemm_sp.cc:
- Use MakeAccessPtrFromRegion for A, B, C, E buffer access pointers
instead of whole-buffer access_ptr. This preserves stage-specific
region offsets from pipeline multi-versioning.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix WS conditional body unwrapping for complex loop bodies
Only unwrap IfThenElse wrapper when the then-branch is a simple flat
sequence of tile-op Evaluate calls. Skip unwrapping for complex bodies
with LetStmt, For, or other control flow that could break variable
scoping when split into producer/consumer for WS.
Fixes variable-used-before-definition error in
blocksparse_attention sparse_gqa_decode_varlen_indice, which has a
conditional loop body containing LetStmt bindings.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Revert num_stages >= 2 to >= 1 for WS candidacy
The >= 2 threshold broke test_num_stages_one_pure_tma_keeps_auto_warp_specialize.
Pure TMA kernels with num_stages=1 should still be WS candidates.
The seer_attention issue (num_stages=1 with manual layout) is handled
by the has_manual_layout_ check, not the num_stages threshold.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Propagate manual layouts onto MVB-expanded buffers via Layout::Expand
When LayoutInference encounters an MVB-expanded buffer (with leading
stage dimensions) whose trailing dimensions match the original layout,
use Layout::Expand to propagate the manual layout instead of rejecting
or skipping the buffer. Applied to all 3 layout propagation paths:
annotated layout map, alias propagation, and finalization.
Also remove the blanket !has_manual_layout_ WS candidate rejection
since manual layouts now survive onto versioned shared buffers via
Layout::Expand.
Fixes test_sparse_ws_regular_metadata_copy_stays_in_producer.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Extend conditional body unwrapping for LetStmt chains and restore
has_manual_layout_ guard
1. LetStmt chain peeling: when IfThenElse then_case starts with
LetStmt bindings, peel them and append to let_bindings before
checking the simple-body guard. This allows WS for conditional
bodies with variable definitions (e.g., sparse attention patterns).
2. Restore !has_manual_layout_ in WS candidacy check: removing it
caused dequant_groupedgemm_bf16_mxfp4_hopper to fail because
MXFP4 layouts don't survive MVB expansion. The Layout::Expand fix
handles sparse metadata layouts but not all manual layout types.
3. Layout::Expand propagation (from Round 6) remains in place for
future use when MVB learns to handle all manual layout types.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Remove is_simple_body guard from conditional WS unwrapping
The simple-body guard only accepted flat Evaluate sequences inside
IfThenElse, blocking legitimate WS for complex conditional bodies
like sparse flash attention (T.clear, T.reduce_max, T.gemm inside
the guard). The LetStmt peeling already handles variable scoping.
Fixes test_pure_tma_consumer_local_init_does_not_leak_into_producer
and test_sparse_ws_regular_metadata_copy_stays_in_producer.
The remaining test_mixed_tma_cp_async_shared_stage_barriers failure
is a pre-existing issue on the original branch: the tiled WS pass
produces SIMT copies instead of cp.async because LowerPTXAsyncCopy
is not in the pass pipeline (the comment at phase.py:308 says it
runs earlier but no actual call exists on either branch or reference).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Preserve pipeline context for WS producer cp.async injection
Wrap the WS producer loop body in a kPipelineContextNumStages AttrStmt
so that LowerTileOp's pipelined_depth_ is > 0 when processing SIMT
producer copies. This enables InjectPTXAsyncCopy to generate cp.async
for global-to-shared copies in the WS producer branch.
Without this, the WS rewriter strips all pipeline annotations from
the rewritten loops, causing LowerTileOp to skip cp.async injection
for SIMT producers. The consumer loop stays annotation-free since it
doesn't need cp.async.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix mixed barrier protocol, consumer-init sinking, restore manual-layout guard
1. Mixed TMA+cp.async barrier: use ptx_cp_async_barrier_noinc for
forward barrier arrival in mixed producer groups, matching the
reference producer_consumer_ws.cc protocol.
2. Consumer-only pre-loop init sinking: in ReplacePipelineLoopInStmt,
guard pre-loop siblings as consumer-only when they're not classified
as producer (TMA/SIMT/cp.async). Fragment init (T.fill, T.clear)
and local buffer init are placed in the consumer branch instead of
the shared prelude.
3. Restored blanket has_manual_layout_ guard: the dtype-based heuristic
to distinguish sparse metadata from MXFP4 layouts doesn't work
because both use uint8. Dequant_groupedgemm_bf16_mxfp4_hopper
requires the guard to prevent broken WS.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Narrow consumer-init sinking, targeted manual-layout check, per-group barrier
1. Consumer-init sinking: only sink pre-loop stmts that are FillNode
writing to fragment/local buffers. Keep block_mask setup and shared
state in the shared prelude.
2. Manual-layout: attempt targeted check that only rejects when TMA
copy destinations match layout_map entries. Collect layout_map vars
and compare against TMA copy destinations.
3. Per-group cp.async barrier: use group-level cp.async flag (single
group for now) instead of function-wide boolean.
The 3 WS issue tests still fail because the layout_map annotation
parsing falls through to the conservative rejection path.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix layout_map parsing with Var→Buffer mapping for manual-layout check
The layout_map annotation uses Map<Var, Layout> before LayoutInference
(not Map<Buffer, Layout>). Parse as Map<ObjectRef, ObjectRef> and
handle both key types: Buffer (post-inference) and Var (pre-inference).
For Var keys, look up the corresponding alloc_buffer by data Var match.
Compare collected manual-layout buffers against pipeline copy
destinations: reject only when a manually-laid-out buffer is also
a producer copy target (TMA/SIMT/cp.async) inside the pipeline.
This allows sparse metadata (E_shared, SIMT-copied) onto the WS path
while rejecting dequant MXFP4 (B_shared, SIMT-copied with swizzle).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix SEGFAULT in manual-layout check, use DetectSwizzleMode
1. Fixed SEGFAULT: removed dangling layout_map_layouts_ vector that
was never populated, causing OOB access. Now stores Buffer+Layout
pairs in layout_map_entries_.
2. Use DetectSwizzleMode to distinguish swizzled layouts (MXFP4,
incompatible with MVB) from non-swizzled (sparse metadata, safe).
Swizzled layouts reject WS candidacy; non-swizzled layouts allow it.
3. Removed debug LOG(WARNING) from hot path.
4. Parse layout_map annotation keys as both Buffer and Var (via
Map<ObjectRef, ObjectRef>), resolving Var keys to alloc_buffers.
Results: sparse metadata WS test PASSES, dequant PASSES (protected).
Down to 2 failures: mixed barrier pattern + consumer init sinking.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Simplify consumer-init sinking to Evaluate-only classification
Sink pre-loop Evaluate nodes classified as kConsumer into consumer
branch. Keep For loops (block_mask setup), producer copies, and other
control flow in the shared prelude. This is simpler and safer than
the FillNode scope check approach.
The test_pure_tma_consumer_local_init test still fails because the
T.fill statements are at a different structural level than the
SeqStmt where the pipeline loop lives. Fixing this requires deeper
IR structure analysis.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Sink consumer-only pre-loop init into WS consumer branch
Extract consumer-only pre-loop Evaluate statements (T.fill on
fragments) from the shared prelude and prepend them to the consumer
branch inside the WS if/else structure. This ensures fragment init
like acc_o, logsum, scores_max appears only in the consumer branch,
not in the shared prelude or producer branch.
Uses a two-pass approach: first ReplacePipelineLoopInStmt extracts
consumer-only stmts into extracted_consumer_init_, then the WS body
is rebuilt with the extracted stmts prepended to the consumer branch.
Fixes test_pure_tma_consumer_local_init_does_not_leak_into_producer.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Finalize pipeline refactor fixes
* Refactor CopyNode stride checks for TMA bulk load/store
- Extracted the global stride validation logic into a new static method `CheckGlobalStrides` for better readability and reusability.
- Updated `CheckBulkLoad` and `CheckBulkStore` methods to utilize the new stride checking function, improving code clarity and maintainability.
- Enhanced documentation for the new method to clarify its purpose and requirements.
* Remove redundant copy direction notes from documentation in `copy.h` for stride checks. This simplifies the comments while maintaining clarity on TMA requirements.
* refactor
* layout related fix
* refactor mbarrier with software pipeline.
* remove legacy pass
* Enhance TMA barrier handling and merging logic
- Updated `CopyToTmaCopyRewriter` to conditionally emit arrive barriers based on the last TMA copy.
- Refactored barrier creation to allow merging of TMA barriers when conditions are met, reducing the number of barriers needed.
- Adjusted consumer and producer logic to accommodate merged barriers, ensuring correct synchronization across TMA copies.
- Improved documentation and comments for clarity on the new barrier merging behavior.
* Update example_mhc_post.py to disable main execution and add test function call
- Commented out the main function call in `example_mhc_post.py`.
- Added a call to `tilelang.disable_cache()` and a new `test(n=4096, h=2560)` function for testing purposes.
* Enhance pipeline barrier handling and TMA copy detection
- Updated `RewritePipelineTmaBarriers` to accept additional parameters for loop variable and stage count, improving barrier synchronization.
- Modified `HandleTileOp` to ensure only user-defined T.copy patterns are marked for pipeline barriers, preventing interference with T.tma_copy.
- Adjusted test cases to enforce specific CUDA compute version requirements for better compatibility and consistency.
* lint fix
* Enhance convolution example and TMA handling
- Added kernel source printing in `example_convolution.py` for debugging.
- Disabled the main function call and invoked `run_regression_perf()` to streamline execution.
- Introduced `MakeTmaLeaderCondition` function in `copy.cc` to improve TMA leader-thread condition handling.
- Updated `LowerBulkCopy` and `LowerBulkCopy1D` methods to utilize the new TMA leader condition for better thread management.
- Enhanced `Conv2DIm2ColOpNode` to support barrier annotations, improving synchronization in TMA operations.
* enhance
* Remove example_dequant_groupedgemm_bf16_mxfp4_hopper.py and clean up phase.py by removing unnecessary whitespace. Enhance TMA handling in producer_consumer_ws_tiled.cc to improve barrier synchronization and streamline TMA copy operations.
* fix
* fix
* Eliminate MVB(barrier_only=true) late fixup from OptimizeForTarget
Move pipeline barrier ownership into InjectSoftwarePipeline: create
pipeline_mbar[num_stages] at final expanded size instead of
pipeline_mbar[1] that required late MVB expansion.
Key changes:
- RewritePipelineTmaBarriers creates barriers at num_stages size
- Barrier indices use FloorMod(loop_var - loop_min, num_stages)
- barrier_init has num_stages entries (one per slot, arrive_count=1)
- CopyToTmaCopyRewriter accepts PrimExpr barrier_id (was int)
- phase.py: remove MVB(barrier_only=True), unify both paths under
PlanAndUpdateBufferAllocationLocation
- Fix pre-existing test expectation for tma_copies annotation
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix Codex review issues: num_stages=1 regression, im2col pipeline barriers, test fixes
1. Fix num_stages=1 regression: use tl_pipelined_num_stages annotation
for barrier sizing instead of max_stage+1. Gate barrier creation on
pipeline_depth > 1 so num_stages=1 kernels don't get multi-versioned
pipeline barriers.
2. Extend RewritePipelineTmaBarriers to handle c2d_im2col: annotate
im2col calls with pipeline barrier in CopyToTmaCopyRewriter. Add
im2col TMA recognition to PipelinePlanning. Fix im2col Lower() to
respect emit_arrive annotation from pipeline barriers.
3. Fix test_simple_pipeline: use annotation-checking approach instead of
structural equality with hardcoded tma_copies annotation that fails
on non-Hopper targets.
4. Remove dead kPipelineMVBStageExpr/kPipelineMVBParityExpr/
kPipelineMVBContextNumStages emission from inject_pipeline.cc since
pipeline barriers are now created at final size.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix depth-1 barrier ownership, add regression tests, remove dead kPipelineMVB*
1. Fix depth-1 ownership gap: always create shared pipeline barrier for
TMA copies even when pipeline_depth=1. This prevents LowerTileOp from
allocating separate per-copy internal barriers, keeping num_stages=1
kernels at pipeline_mbar[1] (single slot).
2. Add checked-in lowering regression tests:
- non-WS num_stages=3 TMA GEMM → asserts pipeline_mbar[3], no fallback
- non-WS num_stages=1 TMA GEMM → asserts pipeline_mbar[1], no multi-slot
- non-WS num_stages=3 im2col → asserts pipeline_mbar[3] feeds tma_load_im2col
3. Remove dead kPipelineMVB* definitions from pipeline_utils.h and
consumption code from multi_version_buffer_rewriter.cc (stacks,
explicit parity/version index logic, attr stripping).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix GemmWMMA.lower() signature and autotuner cache backwards compatibility
[P1] Add missing mbar_phase_expr parameter to GemmWMMA.lower() to match
the interface expected by GemmPy.lower() dispatch. Without this, any
RDNA kernel using the WMMA path would fail with TypeError.
[P2] Make out_idx.json loading optional in autotuner load_from_disk().
Older cache directories don't have this file; fallback to
compile_args.out_idx when the file is absent.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Add .humanize to gitignore
* Fix HIP swizzle codegen and CPU scalar GEMM region handling
[P2] Update codegen_hip.cc to parse tvm_tuple(device_func, panel_size)
format for threadblock_swizzle_pattern, matching the annotation format
now emitted by T.use_swizzle(). Previously expected StringImmNode which
would ICHECK-fail.
[P2] Fix gemm_scalar.py: clear only the output tile region (not the
whole buffer), and use the last two dimensions from regions to handle
rank>2 buffers with leading singleton dimensions.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix HIP codegen to parse tvm_tuple swizzle annotation format
Update codegen_hip.cc threadblock_swizzle_pattern handler to parse
tvm_tuple(device_func, panel_size) format, matching CUDA and cutedsl
codegens. Previously expected StringImmNode which ICHECK-fails with
the current annotation format.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Restore kPipelineMVB* annotations — still needed by WS tiled full MVB
The kPipelineMVBStageExpr/ParityExpr/ContextNumStages annotations and
their MVB consumption code cannot be removed: they are emitted by
EmitImpl for ALL pipelines and consumed by the full MVB call inside
ProducerConsumerWarpSpecializedTiled. Removing them broke non-TMA
pipeline kernels like mhc_post that go through the WS path.
Restores: annotation emission in inject_pipeline.cc, constant
definitions in pipeline_utils.h, and consumption code in
multi_version_buffer_rewriter.cc.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Unify barrier multi-versioning in InjectSoftwarePipeline (Direction A)
Move all pipeline barrier multi-versioning into InjectSoftwarePipeline
via ExpandPipelineBarriers, eliminating the need for late
MVB(barrier_only=true) in OptimizeForTarget.
ExpandPipelineBarriers runs before BuildPipeline and handles:
- ISP-created pipeline_mbar (for non-WS TMA pipelines)
- User-written T.alloc_barrier (for manual WS pipelines like softpipe)
Key design: only barriers with explicit ptx_arrive_barrier calls OR
ISP-created local barriers are expanded. Barriers whose arrival is
managed internally by tile-ops (e.g., tcgen05 MMA arrive) are left
unchanged. This distinguishes pipeline sync barriers from hardware-
managed barriers.
Buffer expansion: barrier[N] -> barrier[N * num_stages]
Index rewriting: barrier[idx] -> barrier[stage_expr * N + idx]
Parity rewriting: user_parity -> (iteration_block + offset) % 2
barrier_init replication: [c0,c1] -> [c0,c1,c0,c1] for num_stages=2
Expanded buffers propagate to outer blocks via pending_buffer_remap_.
barrier_init annotations are replicated in VisitStmt_(BlockNode).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* Fix Python 3.9 compat: replace star expression in index with tuple concat
* fix
* fix
* revert changes
* fix
* enhance pipeline
* enhance pipeline
* implement descriptor reuse pass
* Refactor pipeline management and enhance async copy handling
- Updated `example_mhc_pre.py` to replace the main function call with a test function for better testing flexibility.
- Modified `example_gqa_decode.py` to disable argument parsing and added latency measurement for regression performance.
- Enhanced `copy.cc` and `copy.h` with new checks for pipeline-managed cp.async synchronization and improved async copy handling.
- Updated `inject_pipeline.cc` and `pipeline_planning.cc` to refine the handling of tile operations and global/shared buffer checks.
- Added tests to ensure correct behavior of async pipeline and descriptor allocation reuse.
This commit improves the overall pipeline management and async copy capabilities, ensuring better performance and flexibility in the codebase.
* enhance pipeline
* enhance pipeline
* enhance pipeline
* enhance pipeline
* enhance pipeline
* Revert "enhance pipeline"
This reverts commit
|
||
|
KurisandClaude Opus 4.6
|
6e6295fb35 |
[BugFix] Fix CI failures: clean /tmp on self-hosted runners and skip CuTeDSL alloc_global tests (#2009)
1. Add CI step to clean stale JIT temp files (/tmp/*.so, *.cu, *.cubin, tvm-debug-mode-tempdirs) before tests on self-hosted runners. These files accumulate across CI runs and can fill the disk, causing g++ to be killed (SIGTERM) during JIT compilation. 2. Skip CuTeDSL-incompatible example tests that use alloc_global (flash_decoding, deepseek_mla), since the CuTeDSL wrapper does not yet support alloc_global buffers. Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
||
|
botbw
|
6ca795aba7 |
[example] use alloc_global in split-kv decode kernel (#1991)
[Example] apply alloc_global |
||
|
Lei WangandClaude Opus 4.6
|
ded6a99221 |
[Feature] Add Producer-Consumer Warp Specialization and T.tma_copy() API (#1909)
* [Feature] Add Producer-Consumer Warp Specialization and T.tma_copy() API
This PR introduces Producer-Consumer Warp Specialization for sm90+ TMA
pipelines and a new T.tma_copy() API for explicit mbarrier management.
Key changes:
1. **ProducerConsumerWarpSpecialized pass** (new):
Splits pipelined TMA loops into producer (TMA loads) and consumer
(compute) warp groups with back-pressure barriers for buffer reuse.
Works with num_stages >= 1. Producer warp (128 threads) handles
arrive_expect_tx + tma_load; consumer warps handle compute + arrive
on back-pressure barriers.
2. **T.tma_copy() API** (new):
Fire-and-forget TMA copy with a required `barrier` parameter. Unlike
T.copy() which emits producer+wait pairs, T.tma_copy() emits only
arrive_expect_tx + tma_load. User manages synchronization via
T.mbarrier_wait_parity().
3. **MultiVersionBuffer barrier expansion**:
Extended to handle `shared.barrier` scope buffers, expanding them for
pipelining (1D size multiplication instead of prepending a dimension).
Auto-computes mbarrier parity as `(k // num_stages) % 2`.
4. **Pass ordering** (phase.py):
LowerSharedBarrier now runs after MultiVersionBuffer in the TMA path
so barrier buffers retain their `shared.barrier` scope during expansion.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
* fix for annotate reg alloc
* Enhance debugging and performance measurement in examples
- Added print statements to output kernel source for both forward and backward sparse MLA examples, aiding in debugging.
- Updated `sparse_mla_fwd_pipelined.py` to run performance regression tests and print average latency, improving performance tracking.
- Disabled auto cache in `example_mha_sink_fwd_bhsd.py` for better control during testing.
* fix
* Enhance debugging and performance tracking in examples
- Added print statements to output kernel source in various example scripts, including `example_gqa_sink_bwd_bhsd.py`, `example_blocksparse_gemm.py`, and `example_dequant_groupedgemm_bf16_mxfp4_hopper.py`, aiding in debugging.
- Updated default `window_size` in `example_gqa_sink_bwd_bhsd.py` for improved configuration.
- Disabled auto cache in `example_gemm_schedule.py` to enhance control during testing.
- Introduced a new test for preserving protected auto-injected wait groups in the async copy optimization process.
* clean
* Remove unnecessary print statements for kernel source in example scripts
- Eliminated print statements from `example_gqa_sink_bwd_bhsd.py`, `example_blocksparse_gemm.py`, and `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` to clean up output and improve readability.
- These changes streamline the examples while maintaining functionality.
* clean
* lint fix
* Refactor T.tma_copy() and barrier synchronization in examples and tests
- Updated `example_warp_specialize_gemm_copy_1_gemm_0.py` to use `T.tma_copy()` for optimized memory copying with explicit barrier management.
- Enhanced test cases in `test_tilelang_issue_tma_no_ws.py` and `test_tilelang_language_tma_copy.py` to reflect changes in barrier synchronization, ensuring proper functionality with the new T.tma_copy() API.
- Removed unnecessary print statements and adjusted barrier allocations for clarity and performance.
- Cleaned up the `test_tilelang_transform_inject_tma_barrier.py` file by deleting it, as it was no longer needed.
* Add buffer data mapping for remapped buffers in multi_version_buffer_rewriter
- Introduced a check to ensure that the data variable of a buffer is discoverable, allowing the barrier_init annotation update to find the remapped buffer correctly.
- This change enhances the functionality of the buffer remapping process, ensuring proper synchronization in the transformation pipeline.
* Implement TMA im2col load synchronization with mbarrier allocation
- Added support for allocating mbarriers to synchronize TMA im2col loads, enhancing the pipeline's barrier management.
- Introduced logic to handle mbarrier creation and integration into the TMA copy process, ensuring proper synchronization and data transfer.
- Updated the handling of the mbarrier in the TMA copy statement to improve performance and maintain functionality within the thread-gated block.
* Implement TMA im2col load synchronization with mbarrier allocation
- Added support for allocating mbarriers to synchronize TMA im2col loads, enhancing the pipeline's barrier management.
- Introduced logic to handle mbarrier initialization and integration within the TMA copy process, ensuring proper synchronization during memory operations.
- Updated the handling of the mbarrier in the im2col operation to improve performance and maintain consistency in multi-threaded environments.
* disable tma lower for test_cutedsl_barrier
* Refactor code formatting in test_tilelang_jit_cutedsl.py
- Improved readability by adjusting the formatting of the `tilelang.compile` function call.
- Enabled the main testing function in the script to run without being commented out.
* Enhance sparse_mla_fwd.py and thread_storage_sync.cc with code cleanup and new functionality
- Added a call to disable cache in the main execution block of `sparse_mla_fwd.py` for improved performance.
- Commented out unused benchmarking code in `test_sparse_mla_fwd` to streamline testing.
- Introduced new utility functions in `thread_storage_sync.cc` to analyze and optimize shared memory access patterns, enhancing synchronization management.
- Improved readability and organization of the code by restructuring and commenting on key sections.
* Add thread storage synchronization utilities and optimizations
- Introduced functions to identify shared-like scopes, buffers, and pointer variables, enhancing memory management.
- Implemented a `BarrierTransparentStmtAnalyzer` to analyze statement transparency concerning shared memory.
- Developed a `ThreadStorageSyncOptimizer` to optimize sequences of statements by removing unnecessary synchronization calls, improving performance in multi-threaded environments.
- Enhanced handling of shared memory access in various call types, ensuring proper synchronization during memory operations.
* remove optimize thread_storage_sync
* Refactor pipeline planning to enhance cp.async group scheduling
- Introduced a last-use index for cp.async groups to improve scheduling accuracy.
- Removed unnecessary tracking of first consumers, simplifying the scheduling logic.
- Updated the handling of cp.async and commit statements to ensure proper group ordering based on last-use.
- Added a new test to verify that cp.async groups are ordered correctly by their last-use index.
* lint fix
* Add CUDA requirements to TMA copy tests and update main execution call
- Added decorators to require CUDA and a minimum compute version for the `test_tma_copy_pipeline_2_stages` and `test_tma_copy_pipeline_3_stages` functions.
- Updated the main execution block to call `tilelang.testing.main()` instead of directly invoking the test functions, improving test execution management.
* Refactor attention sink examples by removing pipelined implementations
- Updated imports in `benchmark_gqa_sink_fwd.py` and `benchmark_mha_sink_fwd.py` to use the non-pipelined versions of the attention sink functions.
- Deleted the pipelined implementations of `example_gqa_sink_fwd_bhsd_wgmma_pipelined.py` and `example_mha_sink_fwd_bhsd_wgmma_pipelined.py`.
- Cleaned up regression and test files to remove references to the deleted pipelined examples, ensuring consistency across the codebase.
* performance fix
* Enhance example_blocksparse_gemm.py and inject_pipeline.cc with debugging and utility improvements
- Added a print statement to output the kernel source in `example_blocksparse_gemm.py` for debugging purposes.
- Updated `inject_pipeline.cc` to include a new utility function for stripping TMA copy write buffer attributes, improving code clarity and functionality.
- Modified `producer_consumer_ws.cc` to utilize the new utility function, enhancing the handling of TMA barriers.
- Expanded the `_compile_tvm_ffi` function in the test file to accept additional keyword arguments, improving flexibility for test configurations.
- Introduced a new test to verify that pure TMA warp specialization correctly emits barriers, ensuring compliance with expected behavior.
* lint fix
* Refactor guard variable declarations in producer_consumer_ws.cc for clarity
- Changed guard variable declarations from non-const to const references to improve performance and readability.
- This minor adjustment enhances the handling of producer guards in the context of TMA barriers.
* lint fix
* fix
* lint fix
* perf fix
* Add print statement for kernel source and update main execution in example_dequant_gemm_w4a8.py
- Introduced a print statement to output the kernel source for debugging purposes.
- Updated the main execution block to call `run_regression_perf()` and print the latency, enhancing performance measurement capabilities.
* fix
* Add debugging output and update execution flow in example_tilelang_block_sparse_attn.py
- Added a print statement to display the kernel source for better debugging.
- Updated the main execution block to disable cache and call `run_regression_perf()`, printing the resulting latency for performance measurement.
- Enhanced the producer-consumer logic in producer_consumer_ws.cc with new utility functions to manage local variable access and improve code clarity.
- Introduced a new test case in test_tilelang_issue_tma_no_ws.py to ensure consumer-only local initializations do not leak into the producer, verifying correct behavior in TMA contexts.
* Revert "Add debugging output and update execution flow in example_tilelang_block_sparse_attn.py"
This reverts commit
|
||
|
Lei Wang
|
191d879ddd |
[Bugfix] Fix thread storage sync conflict detection for loop carry write-after-read (#1781)
* Fix thread storage synchronization logic in `thread_storage_sync.cc` to correctly identify conflicts between read and write operations based on loop carry conditions. * lint fix * Refactor `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` to use shared memory for `sorted_token_ids` instead of local memory, improving thread synchronization. Adjust default argument values for M, N, and K in the main function for better testing scenarios. * Add UniformExprChecker to enforce thread synchronization rules Introduce the UniformExprChecker class to determine if expressions are uniform across threads, crucial for safe synchronization in conditional statements. Update the TileLangThreadSyncPlanner to hoist synchronization points out of non-uniform if-statements to prevent potential deadlocks. Enhance tests to validate sync hoisting behavior for various non-uniform conditions involving thread indices and shared memory access. * lint fix * Enhance `example_dequant_groupedgemm_bf16_mxfp4_hopper.py` with cache disabling and kernel source printing for debugging. Update thread synchronization logic in `thread_storage_sync.cc` to check for runtime-dependent conditions, preventing potential deadlocks by hoisting sync points as necessary. * Update submodule `tvm` to latest commit and remove deprecated `example_gqa_decode_varlen_logits_paged.py` file. Refactor `example_gqa_decode_varlen_logits.py` to enhance performance and maintainability by removing unused imports and optimizing shared memory usage. Adjust test cases to reflect the removal of the paged example. * fix * Enhance thread synchronization logic in `thread_storage_sync.cc` by adding a configurable warp size parameter to `RuntimeDependentConditionChecker` and `TileLangThreadSyncPlanner`. This allows for better adaptability to different target architectures. Update the logic to ensure thread extent is a constant and improve handling of runtime-dependent conditions. * lint fix * Refactor thread extent validation in `thread_storage_sync.cc` to use pointer checks instead of optional values. This change improves clarity and ensures that the thread extent is correctly validated as a constant. * Adjust loop variable constraints in `thread_storage_sync.cc` for loop-carry analysis by modifying the extent calculation. This change ensures valid iteration comparisons by reducing the extent by one, allowing for accurate analysis of loop iterations. * lint fix |
||
|
Lei Wang
|
651f885ed7 |
Refactor: Use centralized do_bench from tilelang.profiler (#1670)
* refactor: replace local do_bench function with centralized profiler import - Removed local implementations of the `do_bench` function from multiple example files. - Updated imports to use the centralized `do_bench` function from `tilelang.profiler`, promoting code reuse and consistency across examples. * refactor: remove unused do_bench import from example_chunk_o_bwd.py and __init__.py - Eliminated the unused `do_bench` import from `example_chunk_o_bwd.py` and `tilelang/profiler/__init__.py`, streamlining the code and improving clarity. * lint fix |
||
|
Lei Wang
|
e23bce79d6 |
[Refactor] Simplify do_bench calls by using default warmup and rep parameters (#1568)
* Refactor benchmarking calls to remove warmup parameter in do_bench function Updated multiple examples to streamline benchmarking by removing the warmup parameter from the do_bench function calls. This change simplifies the function signature while maintaining the backend specification for performance measurement. * Update regression performance parameters across multiple examples to support larger input sizes Modified the `run_regression_perf` function in various examples to set default parameters for matrix dimensions (M, N, K) and heads to 4096 or 32, enhancing the benchmarking capabilities for larger datasets. This change aims to improve performance testing and ensure consistency across examples. |
||
|
Gabriel WuandLei Wang
|
0006621b65 |
[Misc] add env for default target/backend/verbose (#1512)
* [Misc] add env for default target/backend/verbose * fix: target_host signature * fix: move all env logic to kernel_cache * fix: example * fix: type hint * Update example_gqa_decode_varlen_logits.py --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com> |
||
|
|
42697c00e6 |
[Cleanup] Remove unnecessary macros in tilelang examples (#1514)
* Remove unnecessary macros in tilelang examples * fix typo --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com> |
||
|
Zhengju Tang
|
feb106b733 | [BugFix] Phaseout unused tests for gqa decode kernels and add the kernels to CI (#1515) | ||
|
Lei Wang
|
09385e7deb |
[Refactor] Support auto swizzling for tma store and phaseout related layout annotations (#1509)
* Remove unnecessary swizzled layout annotations from various attention sink examples and kernel analysis script for improved clarity and performance. * Remove unnecessary swizzled layout annotations from various examples to enhance code clarity and performance. * lint fix * Uncomment main test execution for vectorized cast * Add swizzled layout annotation for B_shared in dequantize GEMM example 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. * Remove unnecessary cache disabling in GQA sink example for improved clarity and performance. * lint fix * Refactor layout functions to support row-major linear layout for any 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. * Refactor `make_linear_layout` to accept a single argument and support 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. * skip callback test * Remove example_dequant_gemm_bf16_mxfp4_hopper_tma.py file, eliminating unused code related to dequantization GEMM example. * typo fix |
||
|
Wenhao Xieandyyttt6
|
ba23181d23 |
[CI] Add preformance regression test script (#1489)
* [Feature]: Add benchmark scripts for examples * apply cupti * fix * format * initial commit * fix * upd * upd * lint * fix * fake * Simplify PR regression test workflow Removed redundant 'Clean pip environment' steps from the workflow. * Update test_perf_regression.py * Enhance regression test bot workflow file handling Updated the GitHub Actions workflow to improve file handling for the regression test report. * Update regression test workflow for artifact naming * Update pr-regression-test-bot.yml * fix * lint * Update performance regression test trigger conditions --------- Co-authored-by: yyttt6 <1652272478@qq.com> |
||
|
Lei Wang
|
c750fb8a98 |
[Enhancement] Update examples and tests for improved type handling functionality (#1448)
* [Enhancement] Update examples and tests for improved type handling and functionality
- Enhanced various example scripts to support new data types and improve compatibility with PyTorch.
- Updated tests across multiple modules to ensure correct functionality with the latest changes in type handling.
- Refactored code in examples to streamline operations and improve clarity, particularly in tensor operations and memory management.
- Added comprehensive tests for new features and fixed existing issues related to type conversions and buffer handling.
* [Refactor] Update accumulation data type to float32 across examples
- Changed accumulation data type from "float" to T.float32 in multiple example scripts to ensure consistency and improve numerical stability.
- This update affects various modules including flash attention, GEMM analysis, convolution, and deepseek MLA examples, enhancing type handling across the board.
* [Refactor] Standardize data type usage across benchmark scripts
- Updated data type definitions in benchmark scripts to use T.float16 and T.float32 consistently, enhancing clarity and type handling.
- Adjusted dtype assignments in matmul functions and configuration setups to align with the new standard.
- Improved overall code consistency and maintainability by ensuring uniform data type usage across various modules.
* [Refactor] Standardize data type usage in templates and scripts
- Updated data type definitions in various templates and scripts to use string representations (e.g., "float16", "int32") instead of T.float16 and T.int32 for improved consistency and clarity.
- Enhanced overall code maintainability by ensuring uniform data type usage across multiple modules, including convolution, elementwise operations, and matrix multiplication templates.
- This change aims to streamline type handling and improve compatibility with existing workflows.
* [Refactor] Standardize data type usage in examples and benchmarks
- Updated data type definitions in various example and benchmark scripts to use T.float16 and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in kernel functions and configuration setups to align with the new standard.
- Improved overall code consistency by ensuring uniform data type usage across multiple modules, including attention mechanisms, matrix multiplication, and GEMM examples.
* [Refactor] Import dtypes from language.v2 module
- Added import statement for dtypes from the language.v2 module to enhance type handling and maintain consistency across the codebase.
- This change aims to streamline data type management and improve overall code clarity.
* fix
* [Refactor] Standardize data type usage across scripts
- Updated data type definitions in various scripts to use string representations (e.g., "float16", "int8") instead of T.float16 and T.int8 for improved consistency and clarity.
- Adjusted dtype assignments in functions and configuration setups to align with the new standard, enhancing overall code maintainability.
- This change affects multiple modules, including benchmark and attention mechanisms, ensuring uniform data type usage throughout the codebase.
* [Refactor] Update data type handling for consistency and clarity
- Changed string representations of data types in the Hint class to use T.float32 and T.int32 for improved consistency.
- Added new data types "int4" and "int16" to the dtypes module, enhancing type support across the codebase.
- Updated function signatures and assertions in the lop3 and mxfp modules to utilize the new data types, ensuring uniformity in type handling.
- This refactor aims to streamline data type management and improve overall code clarity and maintainability.
* [Enhancement] Improve data type handling and error messaging
- Introduced a mapping for canonical data types to their display strings, enhancing clarity in type representation.
- Updated the dtype creation logic to utilize the new mapping, ensuring more intuitive handling of string inputs.
- Refined error messages in the lop3 module to provide clearer feedback on invalid source formats, improving debugging and user experience.
* [Fix] Correct boolean flag in GEMM SP test case
- Updated the boolean flag in the test_gemm_sp_sm90 function to ensure proper functionality in the test case.
- This change enhances the accuracy of the test and aligns it with expected behavior for the GEMM SP implementation.
* [Refactor] Standardize data type usage across scripts
- Updated data type definitions in various scripts to use T.float16 and T.bfloat16 consistently, enhancing clarity and maintainability.
- Adjusted dtype assignments in function signatures and argument parsing to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change affects multiple modules, including benchmarks and examples, improving overall code consistency and readability.
* [Refactor] Standardize data type usage in various modules
- Updated data type assignments in multiple scripts to utilize T.float32, T.int8, and T.int32 consistently, enhancing clarity and maintainability.
- Adjusted function signatures and parameter types across benchmarks, examples, and tests to align with the new standard, ensuring uniform data type usage throughout the codebase.
- This change improves overall code consistency and readability, impacting modules related to matrix multiplication, GEMM, and tensor operations.
* [Refactor] Update argument parsing for data types in benchmarks
- Changed argument parsing for data types in benchmark_matmul_intrinsic.py and benchmark_matmul_sp.py to use string representations ("float16", "int8", "float") instead of T.float16 and T.float.
- This update enhances consistency in data type handling across benchmark scripts, improving clarity and maintainability.
* [Refactor] Update data type handling in benchmark and example scripts
- Changed data type arguments in benchmark and example scripts to use string representations ("float16") instead of T.float16 for improved consistency.
- Updated function signatures and argument parsing to align with the new standard, enhancing clarity and maintainability across the codebase.
- This change affects multiple modules related to attention mechanisms and tensor operations, ensuring uniform data type usage throughout the examples.
* [Refactor] Fix data type conversion in multiple scripts
- Corrected the usage of the data type conversion method from dtype..as_torch() to dtype.as_torch() across various benchmark and example scripts.
- This change enhances consistency in data type handling and improves code readability, impacting modules related to attention mechanisms and tensor operations.
* [Refactor] Update float8 data type usage across multiple scripts
- Changed instances of T.float8_e4m3 to T.float8_e4m3fn in various benchmark, example, and test scripts to ensure consistency in data type handling.
- This update enhances clarity and maintainability across the codebase, particularly in modules related to matrix multiplication and tensor operations.
* [Refactor] Enhance float8 data type handling in CUDA code generation
- Updated the handling of float8 data types in the CUDA code generation to include additional float8 variants, improving type conversion logic.
- Adjusted conditions to ensure proper type checks for float8 conversions, enhancing clarity and maintainability in the codebase.
- Modified layout inference to streamline float8 type checks, ensuring consistency across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.
* [Refactor] Streamline float8 data type handling in CUDA and related modules
- Enhanced float8 data type handling in CUDA code generation by refining type conversion logic and ensuring consistent type checks.
- Updated layout inference for float8 types to improve clarity and maintainability across the implementation.
- This change impacts modules related to matrix operations and CUDA code generation, improving overall type handling and conversion accuracy.
* [Refactor] Remove unnecessary cache disabling in float8 example script
- Eliminated the call to tilelang.disable_cache() in example_group_per_split_token_cast_to_fp8.py to streamline the code.
- This change enhances clarity and maintainability of the example script without affecting its functionality.
* [Refactor] Update data type usage in debug print tests
- Changed the argument for dtype in the test_debug_print_buffer function from a string representation to the corresponding T.bool type.
- This update enhances consistency in data type handling within the test suite, improving clarity and maintainability.
* lint fix
* Update function parameter types from `str` to `T.dtype` for improved type safety in attention sink and related examples
* Refactor `gemv_alloc_reducer` function signature for improved readability by formatting parameters across multiple lines.
|
||
|
Lei Wang
|
29051439db | [Lint] Phaseout Yapf format and embrace ruff format (#1417) | ||
|
Zhengju Tang
|
242b43bbea |
[BugFix] Fix split kernel layout bug of GQA decode (#1386)
* [BugFix] Fix split kernel layout bug of GQA decode * [BugFix] Avoid local with Parallel; use robust fragment instead |
||
|
Tong WUandcoderabbitai[bot]
|
3ab93cd76b |
[Enhancement] Keep max score attention across blocks in FlashAttention for better numerical stablity (#1269)
* Implement max score retention across blocks in FlashAttention for improved stability * fix manual pipeline parameters * Update examples/flash_attention/example_gqa_fwd_varlen.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> * fix typo * more * fix a previous typo --------- Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> |
||
|
Zhengju Tang
|
716dbef52f |
[Example] Add GQA decoding kernel with varlen page table (#1265)
* [Example] Add page table for gqa decode * [Example] Page table for varlen decoding * [Lint] * [Refactor] Remove redundant code * [Lint] * [Lint] * [Lint] |
||
|
Zhengju Tang
|
eb6e89737c |
[GQA] Add varlen decoding kernel with logits saving (#1223)
* [Example] Add GQA varlen decoding kernel with logits return * [Example] Support Sink for GQA varlen decoding * [Example] Add for no-varlen support * [Tune] Add high performance logits saving * [Lint] * [Lint] * [Rename] |
||
|
Yuqi Dong
|
13bdcd605e |
[Refactor]: Change the params in pytest to avoid oom error during ci (#1170)
* [Refactor]: Change the params in pytest to avoid oom error during ci * format * fix * Update test_example_cast.py * Update parameters in test_example_cast * Update test_example_flash_attention.py * update * format * fix * fix * format |
||
|
Tong WU
|
599264cadd |
[Bugfix] Fix CopyNode Lower method to include disable_tma flag in GetCopyInst (#888)
* Fix CopyNode Lower method to include disable_tma flag in GetCopyInst call * Refactor flash attention implementation to disable TMA for specific copy and allow TMA for other operations * attempt to fix lint |
||
|
Tong WU
|
caef45b5a3 |
[Enhancement] Enhance the robustness and generality of MLA examples (#709)
* Enhance format script to automatically install tools in need * Add judgement for small `h_q` in MLA decode examples to improve robustness * Allow scale as a param in MLA decode examples for better generality * Fix typo |
||
|
徐畅
|
98f93db136 | [Bugfix] Remove redundant T.fill to fix precision issue (#667) | ||
|
|
e9a608e222 |
[Bugfix][CI] Bug fixing and migrate CI from ada to hopper (#652)
* fix CI bugs in hopper * lint fix * Update bulk_copy.cc * Refactor bulk copy logic in LowerBulkCopy function - Removed unnecessary blank lines for improved code readability. - Enhanced stride validation by checking for null pointers in global stride calculations, ensuring robustness against symbolic strides. - Updated pass configuration handling in dynamic tile language tests to streamline dynamic alignment and TMA lower pass settings. * test fix * ci fix * Update flash-attention dependencies and clean up example code - Downgraded `flash-attn` dependency version in `requirements-test.txt` to `<=2.2.0`. - Removed unused imports and commented-out code in various example files to enhance readability and maintainability. - Updated the `flashattn` function signature to include default parameters for `block_M`, `block_N`, `num_stages`, and `threads`. - Cleaned up the `example_mha_fwd_varlen.py` and `example_mha_bwd_wgmma_pipelined.py` files by removing unnecessary comments and improving code clarity. - Deleted the `example_mha_inference.py` file as it is no longer needed. * Update CI workflow to remove `--user` flag from pip install commands - Removed the `--user` flag from the pip install commands in both the development and testing sections of the CI workflow to ensure proper installation of dependencies in the virtual environment. * Update CI workflow to include `--no-user` flag in pip install commands - Added the `--no-user` flag to the pip install commands in both the development and testing sections of the CI workflow to ensure dependencies are installed correctly within the virtual environment. * Update CI workflow to include `--no-user` flag in pip install command for wheel mode - Added the `--no-user` flag to the pip install command in the wheel mode section of the CI workflow to ensure dependencies are installed correctly within the virtual environment. * test fix * avoid conflict with system environments * test fix * add commnets --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com> Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Lei Wang
|
d110d0871c |
[Refactor] refactor autotune examples (#617)
* [Refactor] Update tilelang kernel functions and remove unused imports - Refactored the `flashattn_fwd`, `flashattn_bwd_preprocess`, and `flashattn_bwd_postprocess` functions to utilize direct kernel calls instead of cached versions, improving clarity and performance. - Added `@tilelang.jit` decorators with specified output indices to enhance kernel compilation. - Removed unused import of `cached` from `tilelang`, streamlining the code. - Commented out the main testing function call in `test_tilelang_kernel_mha_bwd.py` for potential future use. * [Refactor] Simplify configuration generation in benchmark and example scripts - Refactored the `get_configs` functions in multiple benchmark and example scripts to utilize a dictionary-based approach for parameter configuration, improving readability and maintainability. - Updated the `flashattn` and `chunk_scan_fwd` functions to directly accept configuration parameters, enhancing flexibility in kernel tuning. - Removed redundant code and streamlined the configuration generation process across various files, ensuring consistency in how configurations are defined and utilized. * [Refactor] Update configuration handling in benchmark scripts - Refactored the `get_configs` functions in benchmark scripts to accept a variable argument list, improving flexibility in configuration management. - Enhanced the `matmul` and `flashattn` functions to utilize the updated configuration approach, streamlining parameter handling for kernel tuning. - Added `@autotune` decorators to relevant functions, ensuring consistent autotuning behavior across benchmarks. - Cleaned up redundant code and improved overall readability in the affected files. * [Refactor] Clean up formatting and update subproject commit - Updated the subproject commit reference in the TVM directory to indicate a dirty state. - Removed unnecessary blank lines and improved formatting in the `benchmark_matmul` and `benchmark_matmul_fp8` scripts for better readability. - Streamlined the function definitions in the `flashattn` example script to enhance clarity and maintainability. * [Refactor] Update AutoTuner configuration handling - Modified the AutoTuner class to check if kernel parameters are set before processing tunable arguments, improving robustness in configuration handling. - Enhanced the logic for skipping compilation when tunable parameters are already provided, ensuring efficient use of resources. - Updated comments for clarity and maintainability. * lint fix * Update TVM subproject commit to indicate dirty state and modify MHA backward test cases - Updated the subproject commit reference in the TVM directory to reflect a dirty state. - Adjusted the `test_mha_bwd` function to use a new configuration for the MHA backward tests, changing the context size from 128 to 256. - Uncommented the main testing function call for potential execution. |
||
|
Cunxiao NiandLei Wang
|
3db18726c7 |
[Example] Update examples to use @tilelang.jit (#597)
* [Example] Update kernel compilation in examples to use @tilelang.jit - Refactored multiple examples to eliminate the use of `tilelang.compile` for kernel creation, directly invoking the functions instead. - Added `@tilelang.jit` decorators with appropriate output indices to enhance performance and maintainability. - Improved code clarity by simplifying the kernel invocation process across various examples, ensuring consistency in how kernels are defined and executed. * format * Update example_tilelang_sparse_gqa_decode_varlen_indice.py * Update example_dequant_gemm_fine_grained.py * Update example_gemm_autotune.py --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com> |
||
|
徐畅
|
67d0b6777d |
[BugFix] Fix precision issue in GQA decode when block_N exceeds seqlen/num_split (#575)
* [CI] Add flash_decoding example to CI * Add output of ref latency * format example_gqa_decode.py * [BugFix] Fix precision issue in GQA decode when block_N exceeds seqlen/num_split * format example_gqa_decode.py |
||
|
徐畅
|
7b66fb19c2 |
[CI] Add flash_decoding example to CI (#487)
* [CI] Add flash_decoding example to CI * Add output of ref latency * format example_gqa_decode.py |
||
|
Lei Wang
|
b1ba0cc87f |
[Enhancement] Introduce pass_configs parameter for kernel Caching (#452)
* [Enhancement] Introduce pass_configs parameter for kernel compilation * Added a new `pass_configs` parameter to the `tilelang.compile` function to allow for more flexible kernel compilation configurations. * Updated related classes and methods to accommodate the new parameter, ensuring compatibility across the codebase. * Enhanced the `torch_assert_close` function to include customizable tensor names for better debugging output. * Refactored input handling in example scripts to streamline the process of obtaining inputs for kernel execution. * lint fix |
||
|
yyttt6
|
bfb5b0a3ca |
[Bugfix] fix the unexpected keyword error of autotune (#438)
* yes * [Bugfix] fix the unexpected keyword error of autotune * format * test |
||
|
Lei Wang
|
aa85ddc7dd |
[Doc] Add deepseek_mla to documentation index (#380)
* Add deepseek_mla to documentation index (#380) * lint fix |
||
|
Yuqing XiaandLei Wang
|
8fdfdf03fe |
[Example] Add sparse gqa decode example (#332)
* add example gqa decode wgmma pipelined * add sparse gqa * support num split * support num split * add if condition * add heuristic num split * clean code * add ref * fix bug * add torch ref * fix bug * integrate to torch * symbolic * clean mask * rm actual_num_blocks * clean code * get num_sm via torch * add sparse gqa decode example * format * rm example_gqa_decode_wgmma_pipelined.py * Add license headers to example scripts * format * Remove commented-out cache disabling lines --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com> |
||
|
Lei Wang
|
32060ecda2 |
[Enhancement] Improve flashattn function in example_gqa_decode.py (#329)
- Added a manual seed for reproducibility in PyTorch. - Refactored local variable allocations for better memory management. - Enhanced parallel processing in the flashattn function to improve performance. - Updated layout annotations for clarity and efficiency. These changes optimize the flash attention mechanism and ensure consistent behavior across runs. |
||
|
Zhengju TangandLei Wang
|
5ee58ec737 |
[Dynamic Symbolic] Adaptively vectorize with different condition expressions (#326)
* [Dynamic Symbolic] Adaptively vectorize with different condition expressions * Format * Format * Format * Format * Add MIT License headers to Python files * Simplify return statement in loop vectorization --------- Co-authored-by: Lei Wang <34334180+LeiWang1999@users.noreply.github.com> |
||
|
Lei Wang
|
0fd82ed51f |
[Bugfix] Fix layout conflict issue for gqa decoding examples (#314)
* Remove logging statement from LoopVectorizerDynamic Substitute method for cleaner output. * Refactor flashattn example to improve CUDA configuration handling - Updated the `flashattn` function in `example_gqa_decode.py` to utilize a heuristic configuration based on CUDA device capabilities, enhancing compatibility with different architectures. - Replaced local variable allocations with more efficient constructs and removed unnecessary logging statements for cleaner output. - Adjusted the `do_bench` method call to streamline performance profiling. * lint fix |
||
|
Lei Wang
|
66c7f6a13e |
[Bugfix] Updated autotune usage in the examples to align with the latest changes (#309)
* [Enhancement] Add support for CUDA architecture 8.9 in GEMM template - Introduced conditional inclusion of "gemm_sm89.h" for CUDA architectures 8.9 and above, enhancing compatibility with newer hardware. - This change ensures that the GEMM template can leverage optimizations specific to the 8.9 architecture, improving performance for users with compatible GPUs. * lintfix * [Refactor] Clean up includes in gemm_sm89.h - Removed duplicate inclusion of "common.h" and added "cuda_fp8.h" for improved clarity and organization. - This change enhances the maintainability of the code by ensuring that header files are included only once and in a logical order. * [Enhancement] Improve KernelCache with in-memory caching and detailed docstrings - Added an in-memory cache to the KernelCache class to enhance performance by reducing disk access. - Updated the __new__ method to initialize the memory cache and added logic to check the cache before loading from disk. - Enhanced docstrings across multiple methods to provide clearer explanations of parameters and return values, improving code readability and maintainability. - Implemented a clear_cache method to clear both in-memory and disk caches, ensuring efficient cache management. * lint fix * typofix * [Refactor] Update matmul and flashattn function calls to return structured results - Modified the matmul and flashattn function calls to return a single object containing latency, configuration, and reference latency, improving code clarity and reducing the number of returned variables. - Updated all relevant instances in benchmark and example scripts to accommodate the new return structure, ensuring consistent usage across the codebase. * lint fix |
||
|
Lei Wang
|
bf8a6fc142 |
[Refactor] Deprecated T.Buffer as arguments and rename related calls into T.Tensor (#281)
* [Refactor] Improve flash attention example and layout comparison logic - Removed unnecessary annotation for `lse_local_split` in the flash attention example to streamline the code. - Updated the handling of `lse_local_split` to utilize parallel processing for better performance. - Refactored kernel compilation and profiling logic to enhance clarity and maintainability in the flash attention example. - Added a condition in `FragmentNode::IsEqual` to handle broadcast cases, improving the robustness of layout comparisons. * lint fix * [Enhancement] Add support for shared memory scope in Fill operation - Introduced handling for `shared.dyn` and `shared` memory scopes in the Fill operation. - Implemented parallel operation and layout inference for improved performance in shared memory scenarios. - Updated thread loop partitioning and vectorization logic to accommodate new memory scope handling. * [Refactor] Remove deprecated decorator and enhance Cython kernel handling - Removed the deprecated decorator from the main module and added a new implementation in the utils module for better organization. - Introduced a pointer map in the Cython kernel adapter to manage pointer arguments, improving runtime shape resolution. - Updated the Cython kernel wrapper to utilize the new pointer map for handling kernel arguments. - Enhanced error checking in the tensor utility functions to ensure static shapes are enforced. - Added a new proxy module for buffer and tensor handling, streamlining the interface for TIR programs. * [Feature] Add matrix multiplication test and kernel implementation - Introduced a new test file `test_tilelang_language_ptr.py` that implements a matrix multiplication function using TileLang's primitives. - The `matmul_test` function defines a kernel for performing tile-level GEMM operations with customizable block sizes and data types. - Added a `run_matmul` function to compile and execute the kernel, along with a test function to validate the implementation. - Updated the `proxy.py` file to enhance type handling for buffer and tensor proxies, ensuring compatibility with TIR programs. - Minor formatting improvements in `deprecated.py` for better readability. * lint fix * [Refactor] Update tensor creation in matrix multiplication test - Replaced `T.Tensor.from_ptr` with `T.make_tensor` in `matmul_test` for improved clarity and consistency. - Updated imports in `__init__.py` to include `make_tensor`. - Added `make_tensor` function in `proxy.py` to streamline tensor creation from pointers. * [Refactor] Update tensor definitions across multiple files - Replaced instances of `T.Tensor` with updated tensor definitions in various benchmark and example files to enhance consistency and clarity. - Adjusted tensor shapes and types in functions related to matrix multiplication, attention mechanisms, and other operations. - Improved documentation in README and example files to reflect changes in tensor usage. * lint fix * [Refactor] Update tensor types in attention and matrix multiplication examples - Replaced instances of `T.Tensor` with `T.SharedTensor` and `T.FragmentTensor` in various attention and matrix multiplication functions to improve consistency and clarity. - Adjusted tensor definitions in benchmark and example files to align with the new tensor types. - Enhanced the overall structure and readability of the code by standardizing tensor usage across multiple files. * lint fix * [Refactor] Update tensor types in GEMM example and test files - Replaced instances of `T.Tensor` with `T.LocalTensor` and `T.Buffer` in the GEMM example and related test functions to improve consistency and clarity. - Enhanced the overall structure of the code by standardizing tensor usage across multiple files, aligning with recent updates in tensor definitions. * [Refactor] Update tensor usage in customize.py - Replaced instances of `T.Tensor` with `T.Buffer` in the `reshape` and `view` functions to enhance consistency with recent tensor definitions. - Improved code clarity by standardizing buffer usage across the file. * [Refactor] Update tensor types in test_tilelang_transform_annotate_device_regions.py - Replaced instances of `T.Tensor` with `T.Buffer` in the `before` and `expected` methods of the `TestAnnotateThreadExtent` and `TestAnnotateDeviceScope` classes to enhance consistency with recent tensor definitions. - Improved code clarity by standardizing buffer usage across the test file. * [Refactor] Update tensor types to SharedBuffer and FragmentBuffer - Replaced instances of `T.SharedTensor` and `T.FragmentTensor` with `T.SharedBuffer` and `T.FragmentBuffer` across multiple benchmark, example, and test files to enhance consistency with recent tensor definitions. - Improved code clarity and structure by standardizing buffer usage in attention and matrix multiplication functions. * [Refactor] Introduce Tensor alias for Buffer in proxy.py - Added a new alias `Tensor` for `Buffer` in `proxy.py` to facilitate JIT compilation, ensuring that inputs and outputs are mapped with `torch.Tensor`. - This change enhances clarity and consistency in tensor usage across the codebase. |
||
|
Lei Wang
|
5f5bf53c82 |
[Refactor] Improve flash attention example and layout comparison logic (#270)
* [Refactor] Improve flash attention example and layout comparison logic - Removed unnecessary annotation for `lse_local_split` in the flash attention example to streamline the code. - Updated the handling of `lse_local_split` to utilize parallel processing for better performance. - Refactored kernel compilation and profiling logic to enhance clarity and maintainability in the flash attention example. - Added a condition in `FragmentNode::IsEqual` to handle broadcast cases, improving the robustness of layout comparisons. * lint fix * [Enhancement] Add support for shared memory scope in Fill operation - Introduced handling for `shared.dyn` and `shared` memory scopes in the Fill operation. - Implemented parallel operation and layout inference for improved performance in shared memory scenarios. - Updated thread loop partitioning and vectorization logic to accommodate new memory scope handling. |
||
|
Chaofan LinandLeiWang1999
|
0430cfe7fb |
[Bugfix] Fix Benchmark/Example Code for Autotuning (#254)
* fix tune args * lint * Refactor gemm example and autotuner logging - Updated `ref_program` in `example_gemm.py` to return the result of matrix multiplication instead of modifying an input parameter. - Changed logging filename in `__init__.py` from 'out.log' to 'autotuner.log' for better clarity. - Modified JIT kernel compilation process to include `out_idx` directly in the adapter creation, enhancing flexibility. - Improved validation of `result_idx` in `BaseKernelAdapter` to ensure it falls within valid bounds. * Refactor `ref_program` in `benchmark_matmul_intrinsic.py` to use the `@` operator for matrix multiplication instead of `torch.matmul`, simplifying the implementation by removing the unused parameter `C`. --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com> |
||
|
Yu Cheng
|
3d7b2dc54c |
[Dev][Doc] Enhance Flash Attention Implementation in GQA Decoding Example and Fix Typo (#139)
- Add non-split flash attention macro for more flexible kernel generation - Implement `main_no_split` function to handle single-split scenarios - Modify kernel selection logic to dynamically choose between split and non-split implementations |
||
|
Lei WangandUbuntu
|
7b74bb0128 |
[JIT] Enhance cython/ctypes wrapper for tma descriptor (#126)
* refactor code * enhance tutorial * Enhance error handling and code generation in CUDA and TileLang components This commit introduces several improvements across multiple files: - Added more informative error messages in GEMM layout checks - Updated CUDA codegen to support more flexible function signature generation - Improved TMA descriptor initialization and kernel dispatch logic - Refined library generation and source code parsing utilities - Enhanced error handling in various adapter and wrapper classes * Add thread tag validation for warp specialization Introduce a ThreadTagChecker to validate that a PrimFunc only uses threadIdx.x before applying warp specialization. This prevents unintended transformations on kernels with complex thread binding and provides a clear warning to users about potential issues with warp specialization. * Update TileLang Profiling and Compilation in Flash Decoding Examples Refactor the profiling and compilation workflow in two flash decoding example scripts: - Replace `tilelang.lower()` and `tilelang.Profiler()` with `tilelang.compile()` - Simplify profiler initialization using `get_profiler()` - Update method calls to use the new profiler and compiled kernel objects - Maintain existing performance benchmarking and validation logic * Refactor and clean up code formatting in TileLang testing and adapter modules This commit includes several code style and formatting improvements: - Adjust whitespace and line breaks in test files - Improve code formatting in CUDA source wrapper and adapter utilities - Enhance readability of function calls and argument handling - Remove unnecessary whitespace and standardize indentation - Simplify function signatures and argument parsing * Refactor CUDA codegen and improve code formatting This commit includes several improvements to CUDA code generation and formatting: - Enhance function signature generation in CodeGenTileLangCUDA - Improve code formatting and readability in CUDA-related files - Simplify parameter handling and type annotations - Clean up whitespace and line breaks in codegen and layout files --------- Co-authored-by: Ubuntu <dlisuser@h100testl730RPS.xu5snccwrbtejcqqalluoku5hb.xx.internal.cloudapp.net> |
||
|
Lei Wang
|
13f4b5c6c7 |
[Example] Update GEMM FP8 Example (#123)
* Add DeepSeek MLA decode example with Flash Attention implementation * Add GEMM SplitK and StreamK example implementations This commit introduces two new example scripts demonstrating advanced GEMM (matrix multiplication) techniques: - `example_tilelang_gemm_splitk.py`: Implements a Split-K GEMM kernel using TileLang - `example_tilelang_gemm_streamk.py`: Implements a Stream-K GEMM kernel using TileLang Both examples showcase different parallel computation strategies for matrix multiplication, with comprehensive testing using PyTorch reference implementations. * Refactor GEMM SplitK and StreamK example implementations Clean up and improve code formatting for the SplitK and StreamK GEMM example scripts: - Remove unused import (Profiler) in splitk example - Simplify line breaks and improve code readability - Standardize indentation and remove unnecessary whitespace - Optimize atomic add and copy operations for better clarity * Add block sparse attention benchmarks for multiple libraries This commit introduces comprehensive block sparse attention benchmarks for different libraries: - TileLang block sparse FMHA implementation - Triton block sparse FMHA implementation - PyTorch reference block sparse FMHA implementation - FlashAttention dense FMHA reference implementation The benchmarks include: - Configurable benchmark parameters (batch size, heads, sequence length, etc.) - Sparse mask generation using top-k and threshold methods - Performance measurement for different sparse attention configurations - Utility functions for mask generation and benchmarking * Refactor block sparse attention benchmarks with code style improvements - Add Ruff linter ignore comments to benchmark files - Improve code formatting and line breaks - Remove unused imports - Standardize print statement formatting - Enhance code readability across multiple library benchmarks * lint fix * Add CUDA atomic operations for BFLOAT16 and update function naming - Implement AtomicAdd functions for BFLOAT16 and BFLOAT16x2 in CUDA common header - Rename existing atomic add functions to use PascalCase (atomicAdd -> AtomicAdd) - Add a new __pack_nv_bfloat162 function for packing BFLOAT16 values - Update kernel and language customization to use new function names - Add return type annotations in profiler module * lint fix * Add example for Group Query Attention (GQA) forward pass using Flash Attention in TileLang This commit introduces a new example script `example_gqa_fwd_bshd.py` that demonstrates: - Group Query Attention (GQA) implementation - Flash Attention forward pass - Performance benchmarking - Configurable parameters for batch, heads, sequence length, and dimension - Autotuning support - Reference implementation comparison * Refactor IR lowering pipeline into modular phases This commit introduces a new module `phase.py` to modularize the IR lowering process by splitting the complex lowering pipeline into two distinct phases: - `LowerAndLegalize`: Handles initial IR legalization and transformation - `OptimizeForTarget`: Applies target-specific optimizations The changes simplify the lowering logic in multiple files by extracting the transformation steps into reusable functions, improving code readability and maintainability. * lintfix * nas kernel * Enhance Native Sparse Attention Examples with Code Improvements and Parameter Updates - Updated example_tilelang_nsa.py and example_triton_nsa.py with code formatting and style improvements - Increased default number of heads and selected blocks in TileLang NSA example - Added Ruff linter ignore comments to reference.py - Standardized function signatures and improved code readability across NSA implementations * Add utility math functions for integer operations - Implement `next_power_of_2()` to calculate the next power of 2 for an integer - Add `cdiv()` function for ceiling division of integers * Add utility math functions for integer operations - Implement `next_power_of_2()` to calculate the next power of 2 for an integer - Add `cdiv()` function for ceiling division of integers * Refactor DeepSeek MLA Decode Example with Enhanced Flash Attention Implementation - Update flash attention kernel to support positional embeddings (PE) - Modify reference implementation to handle PE and group query attention - Increase default batch size and adjust benchmarking parameters - Improve kernel performance and readability - Add einops and torch operations for more flexible tensor manipulation * Update README.md with corrected Flash MLA Decoding example path - Modify the example link for Flash MLA Decoding to point to the correct directory - Ensure accurate navigation to the DeepSeek MLA decoding example |
||
|
Yu Cheng
|
b7ca76f1b7 | [Dev] Update MLA decode kernel (#120) | ||
|
Yu Cheng
|
40faabb1f2 |
[Dev] Add MLA and GQA decode examples (#109)
* [CI][Test] Add test cases for tilelang transform MultiVersionBuffer and WarpSpecialized * Relax the mismatch ratio restrictions in the flash_linear_attention and mha tests * [Dev] Add mha backward example * [Dev] Add mla decode example * bug fix * Add triton impl * Add gqa decode example * [Dev] Add GQA decode example * lint * delete unused triton example * set default profiler to 'auto' |
||
|
Yu Cheng
|
bedab1a0f0 |
[CI][Test] Add test cases for tilelang kernel FlashAttention (#54)
* [Dev] Add FlashDecoding example * [CI][Test] Add test cases for tilelang kernel convolution * [CI][Test] Add test cases for tilelang kernel FlashAttention * Reduce the number of stages to ensure the shared memory allocation is valid * Temporarily remove the dim128 case * lint * update einops in requirements-dev.txt * update einops in requirements-test.txt * remove einops in requirements-dev.txt |
||
|
Lei Wang
|
47ecc791c0 |
[Doc] Remove unnecessary layout annotation (#49)
* [Doc] Update documentation structure and content: add overview section, revise project name, and change theme to Furo * [Feature] Add device-side debug printing functions and integrate into kernel interface * lint fix * remove debug print * implement test for debug * lint fix * add some comments * Enhance fragment design and assert fragment print * enhance debug print * add test for msg * lint fix * format * add flash decoding exmaples * remove comment * test simplified |