update mla
This commit is contained in:
parent
9836c8990e
commit
02ce579dfd
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
File diff suppressed because one or more lines are too long
|
|
@ -0,0 +1,7 @@
|
|||
case,q_lengths,kv_lengths,total_q,total_kv,max_len,match_rate,max_abs,worst_ratio,has_nan,has_inf,required_match,pass
|
||||
single_token,1,1,1,1,1,1.00000000,0.00000000,0.00000000,False,False,1.00,True
|
||||
uneven_tail,65;33,65;33,98,98,65,1.00000000,0.01562500,0.32216495,False,False,1.00,True
|
||||
q_shorter_than_kv,32;17,64;41,49,105,64,1.00000000,0.00781250,0.23584904,False,False,0.99,True
|
||||
mixed_ragged,640;384;256;256,1280;1024;768;512,1536,3584,1280,1.00000000,0.00390625,0.15508685,False,False,0.99,True
|
||||
packed_equal_1024,1024,1024,1024,1024,1024,1.00000000,0.01562500,0.29481131,False,False,0.99,True
|
||||
packed_equal_4096,4096,4096,4096,4096,4096,1.00000000,0.01562500,0.32552084,False,False,0.99,True
|
||||
|
|
|
@ -0,0 +1,7 @@
|
|||
case,q_lengths,kv_lengths,total_q,total_kv,max_len,match_rate,max_abs,worst_ratio,has_nan,has_inf,required_match,pass
|
||||
single_token,1,1,1,1,1,1.00000000,0.00000000,0.00000000,False,False,1.00,True
|
||||
uneven_tail,65;33,65;33,98,98,65,1.00000000,0.01562500,0.32216495,False,False,1.00,True
|
||||
q_shorter_than_kv,32;17,64;41,49,105,64,1.00000000,0.00781250,0.23584904,False,False,0.99,True
|
||||
mixed_ragged,640;384;256;256,1280;1024;768;512,1536,3584,1280,1.00000000,0.00390625,0.15508685,False,False,0.99,True
|
||||
packed_equal_1024,1024,1024,1024,1024,1024,1.00000000,0.01562500,0.29481131,False,False,0.99,True
|
||||
packed_equal_4096,4096,4096,4096,4096,4096,1.00000000,0.01562500,0.32552084,False,False,0.99,True
|
||||
|
|
|
@ -0,0 +1,60 @@
|
|||
Loading tilelang libs from dev root: /data/tilelang-metax/build
|
||||
W0717 10:12:26.508000 103879 site-packages/torch/utils/cpp_extension.py:2527] TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
|
||||
W0717 10:12:26.508000 103879 site-packages/torch/utils/cpp_extension.py:2527] If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'] to specific architectures.
|
||||
W0717 10:12:26.554000 103879 site-packages/torch/utils/cpp_extension.py:2527] TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
|
||||
W0717 10:12:26.554000 103879 site-packages/torch/utils/cpp_extension.py:2527] If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'] to specific architectures.
|
||||
W0717 10:12:26.555000 103879 site-packages/torch/utils/cpp_extension.py:2527] TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
|
||||
W0717 10:12:26.555000 103879 site-packages/torch/utils/cpp_extension.py:2527] If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'] to specific architectures.
|
||||
case,batch,total_q,total_kv,max_len,match,max_abs,worst_ratio,nan,inf,pass
|
||||
2026-07-17 10:12:16 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_pv_fullcol_v14` with `out_idx=None`
|
||||
2026-07-17 10:12:20 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:12:26 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_pv_fullcol_v14`
|
||||
1,33,16294,16294,987,1.00000000,0.01562500,0.32552084,False,False,True
|
||||
2026-07-17 10:12:26 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_pv_fullcol_v14` with `out_idx=None`
|
||||
2026-07-17 10:12:29 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:12:35 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_pv_fullcol_v14`
|
||||
2,1,1024,1024,1024,1.00000000,0.01562500,0.32216495,False,False,True
|
||||
2026-07-17 10:12:35 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_pv_fullcol_v14` with `out_idx=None`
|
||||
Traceback (most recent call last):
|
||||
File "/data/operator_task_package/flashinfer_task_package/kernel_ops/FlashinferRaggedPrefill/tests/test_opt_012_formal.py", line 108, in <module>
|
||||
main()
|
||||
File "/data/operator_task_package/flashinfer_task_package/kernel_ops/FlashinferRaggedPrefill/tests/test_opt_012_formal.py", line 63, in main
|
||||
module.run_kernel(
|
||||
File "/data/operator_task_package/flashinfer_task_package/kernel_ops/FlashinferRaggedPrefill/tilelang/opt_014_pv_fullcol.py", line 712, in run_kernel
|
||||
_kernel_cache[key] = build_packed_kernel_pv_fullcol_v14(
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "/data/tilelang-metax/tilelang/jit/__init__.py", line 465, in __call__
|
||||
kernel = self.compile(*args, **kwargs)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "/data/tilelang-metax/tilelang/jit/__init__.py", line 395, in compile
|
||||
kernel_result = compile(
|
||||
^^^^^^^^
|
||||
File "/data/tilelang-metax/tilelang/jit/__init__.py", line 115, in compile
|
||||
return cached(
|
||||
^^^^^^^
|
||||
File "/data/tilelang-metax/tilelang/cache/__init__.py", line 74, in cached
|
||||
return _dispatch_map[execution_backend].cached(
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "/data/tilelang-metax/tilelang/cache/kernel_cache.py", line 345, in cached
|
||||
kernel = JITKernel(
|
||||
^^^^^^^^^^
|
||||
File "/data/tilelang-metax/tilelang/jit/kernel.py", line 136, in __init__
|
||||
adapter = self._compile_and_create_adapter(func, out_idx)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "/data/tilelang-metax/tilelang/jit/kernel.py", line 247, in _compile_and_create_adapter
|
||||
artifact = tilelang.lower(
|
||||
^^^^^^^^^^^^^^^
|
||||
File "/data/tilelang-metax/tilelang/engine/lower.py", line 362, in lower
|
||||
mod = LowerAndLegalize(mod, target)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "/data/tilelang-metax/tilelang/engine/phase.py", line 192, in LowerAndLegalize
|
||||
mod = tilelang.transform.LayoutInference()(mod)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "/data/tilelang-metax/3rdparty/tvm/python/tvm/ir/transform.py", line 167, in __call__
|
||||
return _ffi_transform_api.RunPass(self, mod)
|
||||
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
||||
File "python/tvm_ffi/cython/function.pxi", line 929, in tvm_ffi.core.Function.__call__
|
||||
tvm.error.InternalError: Layout infer conflict between scores and probs in T.Parallel loop:
|
||||
loop Fragment([128, 64] -> [128], replicate: 8, thread: 512, forward_thread: _rep * 64 + _j % 16 // 4 * 16 + _i % 16, forward_index: [_j // 16 * 32 + _i // 16 * 4 + _j % 4], thread_range: I.Range(0, 512))
|
||||
fragment Fragment([128, 64] -> [16], replicate: 1, thread: 512, forward_thread: _i // 16 * 64 + _j % 16 // 4 * 16 + _i % 16, forward_index: [_j // 16 * 4 + _j % 4], thread_range: I.Range(0, 512))
|
||||
|
||||
|
|
@ -0,0 +1,59 @@
|
|||
Loading tilelang libs from dev root: /data/tilelang-metax/build
|
||||
W0717 10:15:35.240000 104250 site-packages/torch/utils/cpp_extension.py:2527] TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
|
||||
W0717 10:15:35.240000 104250 site-packages/torch/utils/cpp_extension.py:2527] If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'] to specific architectures.
|
||||
W0717 10:15:35.286000 104250 site-packages/torch/utils/cpp_extension.py:2527] TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
|
||||
W0717 10:15:35.286000 104250 site-packages/torch/utils/cpp_extension.py:2527] If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'] to specific architectures.
|
||||
W0717 10:15:35.287000 104250 site-packages/torch/utils/cpp_extension.py:2527] TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
|
||||
W0717 10:15:35.287000 104250 site-packages/torch/utils/cpp_extension.py:2527] If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'] to specific architectures.
|
||||
case,batch,total_q,total_kv,max_len,match,max_abs,worst_ratio,nan,inf,pass
|
||||
2026-07-17 10:15:25 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:15:28 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:15:34 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
1,33,16294,16294,987,1.00000000,0.01562500,0.32552084,False,False,True
|
||||
2026-07-17 10:15:35 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:15:38 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:15:43 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
2,1,1024,1024,1024,1.00000000,0.01562500,0.32216495,False,False,True
|
||||
2026-07-17 10:15:44 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:15:47 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:15:52 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
3,1,4096,4096,4096,1.00000000,0.01562500,0.32552084,False,False,True
|
||||
2026-07-17 10:15:53 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:15:56 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:16:01 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
4,1,16384,16384,16384,1.00000000,0.01562500,0.31887755,False,False,True
|
||||
2026-07-17 10:16:01 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:16:05 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:16:10 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
5,4,4096,4096,1024,1.00000000,0.01562500,0.32552084,False,False,True
|
||||
2026-07-17 10:16:10 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:16:13 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:16:19 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
6,4,16384,16384,4096,1.00000000,0.01562500,0.32216495,False,False,True
|
||||
2026-07-17 10:16:19 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:16:22 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:16:28 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
7,16,16384,16384,1024,1.00000000,0.01562500,0.32552084,False,False,True
|
||||
2026-07-17 10:16:28 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:16:31 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:16:37 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
8,16,32768,32768,2048,1.00000000,0.01562500,0.32552084,False,False,True
|
||||
2026-07-17 10:16:37 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:16:41 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:16:47 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
9,4,2048,4096,1024,1.00000000,0.00195312,0.09765624,False,False,True
|
||||
2026-07-17 10:16:47 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:16:50 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:16:56 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
10,4,1536,3584,1280,1.00000000,0.00195312,0.09765624,False,False,True
|
||||
2026-07-17 10:16:57 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:17:00 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:17:06 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
11,2,1024,3072,2048,1.00000000,0.00195312,0.09765624,False,False,True
|
||||
2026-07-17 10:17:07 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:133): TileLang begins to compile kernel `packed_kernel_v_prefetch_after_qk_v15` with `out_idx=None`
|
||||
2026-07-17 10:17:10 [TileLang:tilelang:WARNING] (phase.py:274): Should support PDL for maca target
|
||||
2026-07-17 10:17:16 [TileLang:tilelang.jit.kernel:INFO] (kernel.py:141): TileLang completes to compile kernel `packed_kernel_v_prefetch_after_qk_v15`
|
||||
12,27,12251,12251,873,1.00000000,0.01562500,0.32552084,False,False,True
|
||||
13,15,969,969,123,1.00000000,0.01562500,0.32552084,False,False,True
|
||||
14,1,1,1,1,1.00000000,0.00000000,0.00000000,False,False,True
|
||||
15,2,98,98,65,1.00000000,0.01562500,0.32216495,False,False,True
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
case_id,config,batch,total_q,total_kv,max_q,max_kv,seq_len,tk_ms,tb_ms,speedup,score_ratio,display_score,pass
|
||||
1,ragged_b33_16294,33,16294,16294,987,987,987,2.585,1.575,0.609,0.352,35,True
|
||||
2,equal_b1_s1024,1,1024,1024,1024,1024,1024,0.291,0.260,0.893,0.468,46,True
|
||||
3,equal_b1_s4096,1,4096,4096,4096,4096,4096,3.015,1.655,0.549,0.307,30,True
|
||||
4,equal_b1_s16384,1,16384,16384,16384,16384,16384,45.911,22.426,0.488,0.265,26,True
|
||||
5,equal_b4_s1024,4,4096,4096,1024,1024,1024,0.904,0.635,0.702,0.393,39,True
|
||||
6,equal_b4_s4096,4,16384,16384,4096,4096,4096,11.688,6.064,0.519,0.286,28,True
|
||||
7,equal_b16_s1024,16,16384,16384,1024,1024,1024,3.335,2.024,0.607,0.344,34,True
|
||||
8,equal_b16_s2048,16,32768,32768,2048,2048,2048,12.394,6.638,0.536,0.301,30,True
|
||||
9,q512_k1024_b4,4,2048,4096,512,1024,1024,0.715,0.538,0.752,0.415,41,True
|
||||
10,mixed_b4,4,1536,3584,640,1280,1280,0.577,0.412,0.714,0.400,40,True
|
||||
11,q_lt_kv_b2,2,1024,3072,512,2048,2048,0.701,0.410,0.585,0.344,34,True
|
||||
12,ragged_b27_12251,27,12251,12251,873,873,873,1.816,1.169,0.644,0.369,37,True
|
||||
13,short_ragged_969,15,969,969,123,123,123,0.308,0.151,0.490,0.321,32,True
|
||||
14,single_token,1,1,1,1,1,1,0.018,0.106,5.889,0.855,85,True
|
||||
15,tail_non_power2,2,98,98,65,65,65,0.053,0.109,2.057,0.675,67,True
|
||||
|
|
|
@ -0,0 +1,66 @@
|
|||
# TileLang opt_012 OJ Results
|
||||
|
||||
**Date:** 2026-07-16
|
||||
**File:** `tilelang/opt_012_dense_softmax_cleanup.py`
|
||||
**Parent:** `opt_011_reverse_dense_tiles.py`
|
||||
|
||||
## Summary
|
||||
|
||||
All 15 XPU-OJ testcases passed correctness. This version refines the online
|
||||
softmax computation path for improved numerical efficiency.
|
||||
|
||||
## Results
|
||||
|
||||
| Case | Config | User (ms) | Baseline (ms) | Speedup | Score % | Display |
|
||||
|------|--------|-----------|---------------|---------|---------|---------|
|
||||
| 1 | ragged_b33_16294 | 2.585 | 1.575 | 0.609x | 35.21% | 35 |
|
||||
| 2 | equal_b1_s1024 | 0.291 | 0.260 | 0.893x | 46.77% | 46 |
|
||||
| 3 | equal_b1_s4096 | 3.015 | 1.655 | 0.549x | 30.71% | 30 |
|
||||
| 4 | equal_b1_s16384 | 45.911 | 22.426 | 0.488x | 26.52% | 26 |
|
||||
| 5 | equal_b4_s1024 | 0.904 | 0.635 | 0.702x | 39.26% | 39 |
|
||||
| 6 | equal_b4_s4096 | 11.688 | 6.064 | 0.519x | 28.65% | 28 |
|
||||
| 7 | equal_b16_s1024 | 3.335 | 2.024 | 0.607x | 34.44% | 34 |
|
||||
| 8 | equal_b16_s2048 | 12.394 | 6.638 | 0.536x | 30.08% | 30 |
|
||||
| 9 | q512_k1024_b4 | 0.715 | 0.538 | 0.752x | 41.47% | 41 |
|
||||
| 10 | mixed_b4 | 0.577 | 0.412 | 0.714x | 39.95% | 40 |
|
||||
| 11 | q_lt_kv_b2 | 0.701 | 0.410 | 0.585x | 34.39% | 34 |
|
||||
| 12 | ragged_b27_12251 | 1.816 | 1.169 | 0.644x | 36.93% | 37 |
|
||||
| 13 | short_ragged_969 | 0.308 | 0.151 | 0.490x | 32.12% | 32 |
|
||||
| 14 | single_token | 0.018 | 0.106 | 5.889x | 85.49% | 85 |
|
||||
| 15 | tail_non_power2 | 0.053 | 0.109 | 2.057x | 67.50% | 67 |
|
||||
|
||||
## Aggregate Metrics
|
||||
|
||||
| Metric | Value |
|
||||
|--------|-------|
|
||||
| Sum of user times | 84.311 ms |
|
||||
| Sum of baseline times | 44.172 ms |
|
||||
| Mean score ratio | 40.633% |
|
||||
| Mean display score | 40.3 / 100 |
|
||||
|
||||
These aggregate values were recomputed directly from all 15 rows in the CSV.
|
||||
|
||||
## Key Observations
|
||||
|
||||
1. **Correctness:** All 15 testcases passed
|
||||
2. **Edge cases excel:** Cases 14 and 15 significantly outperform baseline (5.9x and 2.1x)
|
||||
3. **Long dense sequences lag:** Cases 3, 4, 6, 8 show largest gap vs baseline
|
||||
4. **Varlen cases moderate:** Cases 9-13 show reasonable 0.49-0.75x speedup
|
||||
|
||||
## Comparison with opt_011
|
||||
|
||||
Note: opt_011 was tested on 16G machine, opt_012 on 64G machine. Direct
|
||||
comparison may not be fully accurate.
|
||||
|
||||
The performance profile is similar between versions, with the main bottleneck
|
||||
remaining the long dense sequence path where FlashInfer's more sophisticated
|
||||
implementation (partition-KV, async prefetch, MMA-specific optimizations)
|
||||
outperforms the current TileLang mapping.
|
||||
|
||||
## Next Steps
|
||||
|
||||
To close the gap with baseline, future optimization should focus on:
|
||||
1. Improving the long dense sequence path (cases 3, 4, 6, 8)
|
||||
2. Better utilization of C500 matrix instructions
|
||||
3. Potential async prefetch for K/V data
|
||||
4. Consideration of partition-KV for very long sequences
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
##############################
|
||||
Sub-module: Summary
|
||||
------------------------------
|
||||
Name: Total Cycles
|
||||
Description: cycles use by kernel
|
||||
Value: 52,104.54(Kcycles)
|
||||
------------------------------
|
||||
Name: AP busy Duty
|
||||
Description: average AP busy duty of total cycles
|
||||
Value: 99.34%
|
||||
##############################
|
||||
Sub-module: ISU Statistics
|
||||
------------------------------
|
||||
Name: ISU stall cycles layout
|
||||
Description: ISU stall cycles layout
|
||||
Value: {
|
||||
"data": {
|
||||
"wsm_stall": 1230986240.0,
|
||||
"vls_pipeline_stall": 149971994.25641027,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717095614/ISU_stall_cycles_layout20260717095702187.png"
|
||||
}
|
||||
##############################
|
||||
Sub-module: GPU Throughput Statistics
|
||||
------------------------------
|
||||
Name: AP MTE Duty ratio
|
||||
Description: MTE Duty ratio relative to AP active
|
||||
Value: 38.31%
|
||||
------------------------------
|
||||
Name: AP STE Duty ratio
|
||||
Description: STE Duty ratio relative to AP active
|
||||
Value: 1.09%
|
||||
------------------------------
|
||||
Name: AP MMA Duty ratio
|
||||
Description: MMA Duty ratio relative to AP active
|
||||
Value: 20.10%
|
||||
------------------------------
|
||||
Name: VLS Duty ratio
|
||||
Description: VLS Duty ratio relative to AP active
|
||||
Value: 0.0%
|
||||
------------------------------
|
||||
Name: L2C Duty ratio
|
||||
Description: L2C Duty ratio relative to L2C active
|
||||
Value: 8.40%
|
||||
|
|
@ -0,0 +1,17 @@
|
|||
Name,Description,Value
|
||||
Total Cycles,cycles use by kernel,"52,104.54(Kcycles)"
|
||||
AP busy Duty,average AP busy duty of total cycles,99.34%
|
||||
ISU stall cycles layout,ISU stall cycles layout,"{
|
||||
""data"": {
|
||||
""wsm_stall"": 1230986240.0,
|
||||
""vls_pipeline_stall"": 149971994.25641027,
|
||||
""vls_wdata_stall"": 2097152.0,
|
||||
""valu_stall"": 0.0
|
||||
},
|
||||
""filename"": ""/opt/mcProfiler-ubuntu18.04/output20260717095614/ISU_stall_cycles_layout20260717095702187.png""
|
||||
}"
|
||||
AP MTE Duty ratio,MTE Duty ratio relative to AP active,38.31%
|
||||
AP STE Duty ratio,STE Duty ratio relative to AP active,1.09%
|
||||
AP MMA Duty ratio,MMA Duty ratio relative to AP active,20.10%
|
||||
VLS Duty ratio,VLS Duty ratio relative to AP active,0.0%
|
||||
L2C Duty ratio,L2C Duty ratio relative to L2C active,8.40%
|
||||
|
|
|
@ -0,0 +1,56 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"name": "Total Cycles",
|
||||
"description": "cycles use by kernel",
|
||||
"value": "52,104.54(Kcycles)"
|
||||
},
|
||||
{
|
||||
"name": "AP busy Duty",
|
||||
"description": "average AP busy duty of total cycles",
|
||||
"value": "99.34%"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"name": "ISU stall cycles layout",
|
||||
"description": "ISU stall cycles layout",
|
||||
"value": {
|
||||
"data": {
|
||||
"wsm_stall": 1230986240.0,
|
||||
"vls_pipeline_stall": 149971994.25641027,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717095614/ISU_stall_cycles_layout20260717095702187.png"
|
||||
}
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"name": "AP MTE Duty ratio",
|
||||
"description": "MTE Duty ratio relative to AP active",
|
||||
"value": "38.31%"
|
||||
},
|
||||
{
|
||||
"name": "AP STE Duty ratio",
|
||||
"description": "STE Duty ratio relative to AP active",
|
||||
"value": "1.09%"
|
||||
},
|
||||
{
|
||||
"name": "AP MMA Duty ratio",
|
||||
"description": "MMA Duty ratio relative to AP active",
|
||||
"value": "20.10%"
|
||||
},
|
||||
{
|
||||
"name": "VLS Duty ratio",
|
||||
"description": "VLS Duty ratio relative to AP active",
|
||||
"value": "0.0%"
|
||||
},
|
||||
{
|
||||
"name": "L2C Duty ratio",
|
||||
"description": "L2C Duty ratio relative to L2C active",
|
||||
"value": "8.40%"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"data": 52104.538,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Total Cycles"
|
||||
},
|
||||
{
|
||||
"data": 99.33925141030903,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP busy Duty"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"data": {
|
||||
"wsm_stall": 1230986240.0,
|
||||
"vls_pipeline_stall": 149971994.25641027,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "ISU stall cycles layout"
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"data": 38.30926808749678,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP MTE Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 1.0916328894651182,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP STE Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 20.100054369899006,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP MMA Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 0.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "VLS Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 8.399882159167529,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "L2C Duty ratio"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
##############################
|
||||
Sub-module: Summary
|
||||
------------------------------
|
||||
Name: Total Cycles
|
||||
Description: cycles use by kernel
|
||||
Value: 52,021.89(Kcycles)
|
||||
------------------------------
|
||||
Name: AP busy Duty
|
||||
Description: average AP busy duty of total cycles
|
||||
Value: 99.55%
|
||||
##############################
|
||||
Sub-module: ISU Statistics
|
||||
------------------------------
|
||||
Name: ISU stall cycles layout
|
||||
Description: ISU stall cycles layout
|
||||
Value: {
|
||||
"data": {
|
||||
"wsm_stall": 1231045529.6,
|
||||
"vls_pipeline_stall": 149665424.41025642,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717095614/ISU_stall_cycles_layout20260717095702519.png"
|
||||
}
|
||||
##############################
|
||||
Sub-module: GPU Throughput Statistics
|
||||
------------------------------
|
||||
Name: AP MTE Duty ratio
|
||||
Description: MTE Duty ratio relative to AP active
|
||||
Value: 38.30%
|
||||
------------------------------
|
||||
Name: AP STE Duty ratio
|
||||
Description: STE Duty ratio relative to AP active
|
||||
Value: 1.09%
|
||||
------------------------------
|
||||
Name: AP MMA Duty ratio
|
||||
Description: MMA Duty ratio relative to AP active
|
||||
Value: 20.10%
|
||||
------------------------------
|
||||
Name: VLS Duty ratio
|
||||
Description: VLS Duty ratio relative to AP active
|
||||
Value: 0.0%
|
||||
------------------------------
|
||||
Name: L2C Duty ratio
|
||||
Description: L2C Duty ratio relative to L2C active
|
||||
Value: 8.41%
|
||||
|
|
@ -0,0 +1,17 @@
|
|||
Name,Description,Value
|
||||
Total Cycles,cycles use by kernel,"52,021.89(Kcycles)"
|
||||
AP busy Duty,average AP busy duty of total cycles,99.55%
|
||||
ISU stall cycles layout,ISU stall cycles layout,"{
|
||||
""data"": {
|
||||
""wsm_stall"": 1231045529.6,
|
||||
""vls_pipeline_stall"": 149665424.41025642,
|
||||
""vls_wdata_stall"": 2097152.0,
|
||||
""valu_stall"": 0.0
|
||||
},
|
||||
""filename"": ""/opt/mcProfiler-ubuntu18.04/output20260717095614/ISU_stall_cycles_layout20260717095702519.png""
|
||||
}"
|
||||
AP MTE Duty ratio,MTE Duty ratio relative to AP active,38.30%
|
||||
AP STE Duty ratio,STE Duty ratio relative to AP active,1.09%
|
||||
AP MMA Duty ratio,MMA Duty ratio relative to AP active,20.10%
|
||||
VLS Duty ratio,VLS Duty ratio relative to AP active,0.0%
|
||||
L2C Duty ratio,L2C Duty ratio relative to L2C active,8.41%
|
||||
|
|
|
@ -0,0 +1,56 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"name": "Total Cycles",
|
||||
"description": "cycles use by kernel",
|
||||
"value": "52,021.89(Kcycles)"
|
||||
},
|
||||
{
|
||||
"name": "AP busy Duty",
|
||||
"description": "average AP busy duty of total cycles",
|
||||
"value": "99.55%"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"name": "ISU stall cycles layout",
|
||||
"description": "ISU stall cycles layout",
|
||||
"value": {
|
||||
"data": {
|
||||
"wsm_stall": 1231045529.6,
|
||||
"vls_pipeline_stall": 149665424.41025642,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717095614/ISU_stall_cycles_layout20260717095702519.png"
|
||||
}
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"name": "AP MTE Duty ratio",
|
||||
"description": "MTE Duty ratio relative to AP active",
|
||||
"value": "38.30%"
|
||||
},
|
||||
{
|
||||
"name": "AP STE Duty ratio",
|
||||
"description": "STE Duty ratio relative to AP active",
|
||||
"value": "1.09%"
|
||||
},
|
||||
{
|
||||
"name": "AP MMA Duty ratio",
|
||||
"description": "MMA Duty ratio relative to AP active",
|
||||
"value": "20.10%"
|
||||
},
|
||||
{
|
||||
"name": "VLS Duty ratio",
|
||||
"description": "VLS Duty ratio relative to AP active",
|
||||
"value": "0.0%"
|
||||
},
|
||||
{
|
||||
"name": "L2C Duty ratio",
|
||||
"description": "L2C Duty ratio relative to L2C active",
|
||||
"value": "8.41%"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"data": 52021.891,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Total Cycles"
|
||||
},
|
||||
{
|
||||
"data": 99.54653897529407,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP busy Duty"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"data": {
|
||||
"wsm_stall": 1231045529.6,
|
||||
"vls_pipeline_stall": 149665424.41025642,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "ISU stall cycles layout"
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"data": 38.30305307273055,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP MTE Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 1.0910904246074884,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP STE Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 20.0979755297533,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP MMA Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 0.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "VLS Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 8.40757788721454,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "L2C Duty ratio"
|
||||
}
|
||||
]
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 40 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
|
|
@ -0,0 +1,25 @@
|
|||
{
|
||||
"CE": [
|
||||
"CE_PERF_BUSY_CYCLES"
|
||||
],
|
||||
"ISU": [
|
||||
"AP_PERF_ISU_ARB_VALU_STALL",
|
||||
"AP_PERF_INST_TRANS_FP16",
|
||||
"AP_PERF_ISU_ARB_WSM_STALL",
|
||||
"AP_PERF_INSTS_SALU",
|
||||
"AP_PERF_INST_TRANS_FP32",
|
||||
"AP_PERF_ISU_ARB_DATA_STALL",
|
||||
"AP_PERF_INST_CYCLES_VALU_4CYCLES",
|
||||
"AP_PERF_ISU_AP_BUSY",
|
||||
"AP_PERF_MTE_PI_SEG_SEL_F64",
|
||||
"AP_PERF_ISU_ARB_VLS_STALL",
|
||||
"AP_PERF_INST_CYCLES_MMA"
|
||||
],
|
||||
"VLS": [
|
||||
"PERFEVENT_VLS_ACTIVE"
|
||||
],
|
||||
"L2C": [
|
||||
"1_b1",
|
||||
"perf_req"
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,134 @@
|
|||
{
|
||||
"perfListPerBlock": {
|
||||
"CE": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
4,
|
||||
"CE_PERF_SEND_RM_WAVES"
|
||||
],
|
||||
[
|
||||
2,
|
||||
"CE_PERF_BUSY_CYCLES"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"ISU": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
1,
|
||||
"AP_PERF_WAVES"
|
||||
],
|
||||
[
|
||||
82,
|
||||
"AP_PERF_ISU_ARB_VALU_STALL"
|
||||
],
|
||||
[
|
||||
41,
|
||||
"AP_PERF_INST_TRANS_FP16"
|
||||
],
|
||||
[
|
||||
79,
|
||||
"AP_PERF_ISU_ARB_WSM_STALL"
|
||||
],
|
||||
[
|
||||
10,
|
||||
"AP_PERF_INSTS_SALU"
|
||||
],
|
||||
[
|
||||
40,
|
||||
"AP_PERF_INST_TRANS_FP32"
|
||||
],
|
||||
[
|
||||
81,
|
||||
"AP_PERF_ISU_ARB_DATA_STALL"
|
||||
]
|
||||
]
|
||||
},
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 1,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
1,
|
||||
"AP_PERF_WAVES"
|
||||
],
|
||||
[
|
||||
52,
|
||||
"AP_PERF_INST_CYCLES_VALU_4CYCLES"
|
||||
],
|
||||
[
|
||||
89,
|
||||
"AP_PERF_ISU_AP_BUSY"
|
||||
],
|
||||
[
|
||||
101,
|
||||
"AP_PERF_MTE_PI_SEG_SEL_F64"
|
||||
],
|
||||
[
|
||||
80,
|
||||
"AP_PERF_ISU_ARB_VLS_STALL"
|
||||
],
|
||||
[
|
||||
43,
|
||||
"AP_PERF_INST_CYCLES_MMA"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"VLS": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
23,
|
||||
"PERFEVENT_VLS_ACTIVE"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"L2C": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
1,
|
||||
"1_b1"
|
||||
],
|
||||
[
|
||||
74,
|
||||
"perf_req"
|
||||
]
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"action": "config,disable,read_dump,enable,kernel_collect",
|
||||
"conditions": {
|
||||
"kernelnames": [
|
||||
"packed_kernel_dense_softmax_cleanup_v12_kernel"
|
||||
],
|
||||
"include": true,
|
||||
"is_perkernel": true,
|
||||
"is_multi_device": false,
|
||||
"counts": 4
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,134 @@
|
|||
{
|
||||
"perfListPerBlock": {
|
||||
"CE": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
4,
|
||||
"CE_PERF_SEND_RM_WAVES"
|
||||
],
|
||||
[
|
||||
2,
|
||||
"CE_PERF_BUSY_CYCLES"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"ISU": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
1,
|
||||
"AP_PERF_WAVES"
|
||||
],
|
||||
[
|
||||
82,
|
||||
"AP_PERF_ISU_ARB_VALU_STALL"
|
||||
],
|
||||
[
|
||||
41,
|
||||
"AP_PERF_INST_TRANS_FP16"
|
||||
],
|
||||
[
|
||||
79,
|
||||
"AP_PERF_ISU_ARB_WSM_STALL"
|
||||
],
|
||||
[
|
||||
10,
|
||||
"AP_PERF_INSTS_SALU"
|
||||
],
|
||||
[
|
||||
40,
|
||||
"AP_PERF_INST_TRANS_FP32"
|
||||
],
|
||||
[
|
||||
81,
|
||||
"AP_PERF_ISU_ARB_DATA_STALL"
|
||||
]
|
||||
]
|
||||
},
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 1,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
1,
|
||||
"AP_PERF_WAVES"
|
||||
],
|
||||
[
|
||||
52,
|
||||
"AP_PERF_INST_CYCLES_VALU_4CYCLES"
|
||||
],
|
||||
[
|
||||
89,
|
||||
"AP_PERF_ISU_AP_BUSY"
|
||||
],
|
||||
[
|
||||
101,
|
||||
"AP_PERF_MTE_PI_SEG_SEL_F64"
|
||||
],
|
||||
[
|
||||
80,
|
||||
"AP_PERF_ISU_ARB_VLS_STALL"
|
||||
],
|
||||
[
|
||||
43,
|
||||
"AP_PERF_INST_CYCLES_MMA"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"VLS": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
23,
|
||||
"PERFEVENT_VLS_ACTIVE"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"L2C": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
1,
|
||||
"1_b1"
|
||||
],
|
||||
[
|
||||
74,
|
||||
"perf_req"
|
||||
]
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"action": "config,disable,read_dump,enable,kernel_collect",
|
||||
"conditions": {
|
||||
"kernelnames": [
|
||||
"packed_kernel_dense_softmax_cleanup_v12_kernel"
|
||||
],
|
||||
"include": true,
|
||||
"is_perkernel": true,
|
||||
"is_multi_device": false,
|
||||
"counts": 4
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,189 @@
|
|||
{
|
||||
"perf_counters": {
|
||||
"0": [
|
||||
{
|
||||
"CE": {
|
||||
"generic": {
|
||||
"CE_PERF_BUSY_CYCLES": 10608946191,
|
||||
"CE_PERF_SEND_RM_WAVES": 1554177
|
||||
}
|
||||
},
|
||||
"ISU": {
|
||||
"dpcId0,blockId0": {
|
||||
"AP_PERF_WAVES": 14976,
|
||||
"AP_PERF_INSTS_SALU": 18112308,
|
||||
"AP_PERF_INST_TRANS_FP32": 85822208,
|
||||
"AP_PERF_INST_TRANS_FP16": 0,
|
||||
"AP_PERF_ISU_ARB_WSM_STALL": 372920866,
|
||||
"AP_PERF_ISU_ARB_DATA_STALL": 700264,
|
||||
"AP_PERF_ISU_ARB_VALU_STALL": 0
|
||||
},
|
||||
"dpcId0,blockId1": {
|
||||
"AP_PERF_WAVES": 15036,
|
||||
"AP_PERF_MTE_PI_SEG_SEL_F64": 0,
|
||||
"AP_PERF_INST_CYCLES_MMA": 1290141696,
|
||||
"AP_PERF_INST_CYCLES_VALU_4CYCLES": 2129860432,
|
||||
"AP_PERF_ISU_ARB_VLS_STALL": 44390461,
|
||||
"AP_PERF_ISU_AP_BUSY": 1605762663
|
||||
}
|
||||
},
|
||||
"L2C": {
|
||||
"dpcId0,blockId0": {
|
||||
"1_b1": 10609002354,
|
||||
"perf_req": 136051869
|
||||
}
|
||||
},
|
||||
"VLS": {
|
||||
"dpcId0,blockId0": {
|
||||
"PERFEVENT_VLS_ACTIVE": 0
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"CE": {
|
||||
"generic": {
|
||||
"CE_PERF_BUSY_CYCLES": 52104538,
|
||||
"CE_PERF_SEND_RM_WAVES": 32768
|
||||
}
|
||||
},
|
||||
"ISU": {
|
||||
"dpcId0,blockId0": {
|
||||
"AP_PERF_WAVES": 320,
|
||||
"AP_PERF_INSTS_SALU": 579520,
|
||||
"AP_PERF_INST_TRANS_FP32": 2764800,
|
||||
"AP_PERF_INST_TRANS_FP16": 0,
|
||||
"AP_PERF_ISU_ARB_WSM_STALL": 12021350,
|
||||
"AP_PERF_ISU_ARB_DATA_STALL": 20480,
|
||||
"AP_PERF_ISU_ARB_VALU_STALL": 0
|
||||
},
|
||||
"dpcId0,blockId1": {
|
||||
"AP_PERF_WAVES": 312,
|
||||
"AP_PERF_MTE_PI_SEG_SEL_F64": 0,
|
||||
"AP_PERF_INST_CYCLES_MMA": 41615360,
|
||||
"AP_PERF_INST_CYCLES_VALU_4CYCLES": 68533184,
|
||||
"AP_PERF_ISU_ARB_VLS_STALL": 1427956,
|
||||
"AP_PERF_ISU_AP_BUSY": 51760258
|
||||
}
|
||||
},
|
||||
"L2C": {
|
||||
"dpcId0,blockId0": {
|
||||
"1_b1": 52083814,
|
||||
"perf_req": 4374979
|
||||
}
|
||||
},
|
||||
"VLS": {
|
||||
"dpcId0,blockId0": {
|
||||
"PERFEVENT_VLS_ACTIVE": 0
|
||||
}
|
||||
},
|
||||
"args": {
|
||||
"kernelName": [
|
||||
"packed_kernel_dense_softmax_cleanup_v12_kernel"
|
||||
],
|
||||
"tag": [
|
||||
-1
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
"CE": {
|
||||
"generic": {
|
||||
"CE_PERF_BUSY_CYCLES": 52021891,
|
||||
"CE_PERF_SEND_RM_WAVES": 32768
|
||||
}
|
||||
},
|
||||
"ISU": {
|
||||
"dpcId0,blockId0": {
|
||||
"AP_PERF_WAVES": 320,
|
||||
"AP_PERF_INSTS_SALU": 579520,
|
||||
"AP_PERF_INST_TRANS_FP32": 2764800,
|
||||
"AP_PERF_INST_TRANS_FP16": 0,
|
||||
"AP_PERF_ISU_ARB_WSM_STALL": 12021929,
|
||||
"AP_PERF_ISU_ARB_DATA_STALL": 20480,
|
||||
"AP_PERF_ISU_ARB_VALU_STALL": 0
|
||||
},
|
||||
"dpcId0,blockId1": {
|
||||
"AP_PERF_WAVES": 312,
|
||||
"AP_PERF_MTE_PI_SEG_SEL_F64": 0,
|
||||
"AP_PERF_INST_CYCLES_MMA": 41631744,
|
||||
"AP_PERF_INST_CYCLES_VALU_4CYCLES": 68559744,
|
||||
"AP_PERF_ISU_ARB_VLS_STALL": 1425037,
|
||||
"AP_PERF_ISU_AP_BUSY": 51785992
|
||||
}
|
||||
},
|
||||
"L2C": {
|
||||
"dpcId0,blockId0": {
|
||||
"1_b1": 52036140,
|
||||
"perf_req": 4374979
|
||||
}
|
||||
},
|
||||
"VLS": {
|
||||
"dpcId0,blockId0": {
|
||||
"PERFEVENT_VLS_ACTIVE": 0
|
||||
}
|
||||
},
|
||||
"args": {
|
||||
"kernelName": [
|
||||
"packed_kernel_dense_softmax_cleanup_v12_kernel"
|
||||
],
|
||||
"tag": [
|
||||
-1
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"umd_data": {},
|
||||
"device_info": {
|
||||
"3": {
|
||||
"gpu_id": "51332",
|
||||
"cpu_cores_count": 0,
|
||||
"mem_banks_count": 1,
|
||||
"direct_link_count": 4,
|
||||
"indirect_link_count": 8,
|
||||
"vendor_id": 39321,
|
||||
"device_id": 16385,
|
||||
"drm_render_minor": 129,
|
||||
"isa_major": 10,
|
||||
"isa_minor": 0,
|
||||
"is_vf": 0,
|
||||
"max_vf_nums": 8,
|
||||
"domain": 0,
|
||||
"location_id": 3840,
|
||||
"num_sdma_engines": 5,
|
||||
"num_sdma_queues_per_engine": 9,
|
||||
"num_vpue_cores": 1,
|
||||
"num_vpud_cores": 8,
|
||||
"caches_count": 157,
|
||||
"peu_id_base": 0,
|
||||
"peu_count": 416,
|
||||
"max_waves_per_peu": 8,
|
||||
"wsm_size_in_kb": 64,
|
||||
"wave_front_size": 64,
|
||||
"dpc_count": 8,
|
||||
"dpc0_ap_mask": 8191,
|
||||
"dpc1_ap_mask": 73727,
|
||||
"dpc2_ap_mask": 139263,
|
||||
"dpc3_ap_mask": 204799,
|
||||
"dpc4_ap_mask": 270335,
|
||||
"dpc5_ap_mask": 335871,
|
||||
"dpc6_ap_mask": 401407,
|
||||
"dpc7_ap_mask": 466943,
|
||||
"dpc_arrays": 1,
|
||||
"ap_per_dpc": 13,
|
||||
"peu_per_ap": 4,
|
||||
"pri_mem_per_thread": 4,
|
||||
"max_slots_private_ap": 32,
|
||||
"num_ce_queues": 16,
|
||||
"max_engine_clk_gpu": 1600,
|
||||
"max_engine_clk_cpu": 3600,
|
||||
"mgpu_id": 0,
|
||||
"topology_id": 2,
|
||||
"socket_id": 3,
|
||||
"hbmecc": 1,
|
||||
"local_mem_size": 68719476736,
|
||||
"capability": 4432512,
|
||||
"maxprocess": 10
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
|
|
@ -0,0 +1,46 @@
|
|||
##############################
|
||||
Sub-module: Summary
|
||||
------------------------------
|
||||
Name: Total Cycles
|
||||
Description: cycles use by kernel
|
||||
Value: 10,608,946.19(Kcycles)
|
||||
------------------------------
|
||||
Name: AP busy Duty
|
||||
Description: average AP busy duty of total cycles
|
||||
Value: 15.14%
|
||||
##############################
|
||||
Sub-module: ISU Statistics
|
||||
------------------------------
|
||||
Name: ISU stall cycles layout
|
||||
Description: ISU stall cycles layout
|
||||
Value: {
|
||||
"data": {
|
||||
"wsm_stall": 38700923661.67748,
|
||||
"vls_pipeline_stall": 4588363494.652634,
|
||||
"vls_wdata_stall": 72671888.53685898,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717095614/ISU_stall_cycles_layout20260717095701860.png"
|
||||
}
|
||||
##############################
|
||||
Sub-module: GPU Throughput Statistics
|
||||
------------------------------
|
||||
Name: AP MTE Duty ratio
|
||||
Description: MTE Duty ratio relative to AP active
|
||||
Value: 38.53%
|
||||
------------------------------
|
||||
Name: AP STE Duty ratio
|
||||
Description: STE Duty ratio relative to AP active
|
||||
Value: 1.13%
|
||||
------------------------------
|
||||
Name: AP MMA Duty ratio
|
||||
Description: MMA Duty ratio relative to AP active
|
||||
Value: 20.09%
|
||||
------------------------------
|
||||
Name: VLS Duty ratio
|
||||
Description: VLS Duty ratio relative to AP active
|
||||
Value: 0.0%
|
||||
------------------------------
|
||||
Name: L2C Duty ratio
|
||||
Description: L2C Duty ratio relative to L2C active
|
||||
Value: 1.28%
|
||||
|
|
@ -0,0 +1,17 @@
|
|||
Name,Description,Value
|
||||
Total Cycles,cycles use by kernel,"10,608,946.19(Kcycles)"
|
||||
AP busy Duty,average AP busy duty of total cycles,15.14%
|
||||
ISU stall cycles layout,ISU stall cycles layout,"{
|
||||
""data"": {
|
||||
""wsm_stall"": 38700923661.67748,
|
||||
""vls_pipeline_stall"": 4588363494.652634,
|
||||
""vls_wdata_stall"": 72671888.53685898,
|
||||
""valu_stall"": 0.0
|
||||
},
|
||||
""filename"": ""/opt/mcProfiler-ubuntu18.04/output20260717095614/ISU_stall_cycles_layout20260717095701860.png""
|
||||
}"
|
||||
AP MTE Duty ratio,MTE Duty ratio relative to AP active,38.53%
|
||||
AP STE Duty ratio,STE Duty ratio relative to AP active,1.13%
|
||||
AP MMA Duty ratio,MMA Duty ratio relative to AP active,20.09%
|
||||
VLS Duty ratio,VLS Duty ratio relative to AP active,0.0%
|
||||
L2C Duty ratio,L2C Duty ratio relative to L2C active,1.28%
|
||||
|
|
|
@ -0,0 +1,56 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"name": "Total Cycles",
|
||||
"description": "cycles use by kernel",
|
||||
"value": "10,608,946.19(Kcycles)"
|
||||
},
|
||||
{
|
||||
"name": "AP busy Duty",
|
||||
"description": "average AP busy duty of total cycles",
|
||||
"value": "15.14%"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"name": "ISU stall cycles layout",
|
||||
"description": "ISU stall cycles layout",
|
||||
"value": {
|
||||
"data": {
|
||||
"wsm_stall": 38700923661.67748,
|
||||
"vls_pipeline_stall": 4588363494.652634,
|
||||
"vls_wdata_stall": 72671888.53685898,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717095614/ISU_stall_cycles_layout20260717095701860.png"
|
||||
}
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"name": "AP MTE Duty ratio",
|
||||
"description": "MTE Duty ratio relative to AP active",
|
||||
"value": "38.53%"
|
||||
},
|
||||
{
|
||||
"name": "AP STE Duty ratio",
|
||||
"description": "STE Duty ratio relative to AP active",
|
||||
"value": "1.13%"
|
||||
},
|
||||
{
|
||||
"name": "AP MMA Duty ratio",
|
||||
"description": "MMA Duty ratio relative to AP active",
|
||||
"value": "20.09%"
|
||||
},
|
||||
{
|
||||
"name": "VLS Duty ratio",
|
||||
"description": "VLS Duty ratio relative to AP active",
|
||||
"value": "0.0%"
|
||||
},
|
||||
{
|
||||
"name": "L2C Duty ratio",
|
||||
"description": "L2C Duty ratio relative to L2C active",
|
||||
"value": "1.28%"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"data": 10608946.191,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Total Cycles"
|
||||
},
|
||||
{
|
||||
"data": 15.135929941488758,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP busy Duty"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"data": {
|
||||
"wsm_stall": 38700923661.67748,
|
||||
"vls_pipeline_stall": 4588363494.652634,
|
||||
"vls_wdata_stall": 72671888.53685898,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "ISU stall cycles layout"
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"data": 38.52569057861815,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP MTE Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 1.132475785840243,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP STE Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 20.086120535236283,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP MMA Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 0.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "VLS Duty ratio"
|
||||
},
|
||||
{
|
||||
"data": 1.2824190669417963,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "L2C Duty ratio"
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
|
|
@ -0,0 +1,64 @@
|
|||
##############################
|
||||
Sub-module: Summary
|
||||
------------------------------
|
||||
Name: Total Cycles
|
||||
Description: cycles use by kernel
|
||||
Value: 52,126.52(Kcycles)
|
||||
------------------------------
|
||||
Name: AP busy Duty
|
||||
Description: average AP busy duty of total cycles
|
||||
Value: 99.20%
|
||||
##############################
|
||||
Sub-module: ISU Statistics
|
||||
------------------------------
|
||||
Name: ISU stall cycles layout
|
||||
Description: ISU stall cycles layout
|
||||
Value: {
|
||||
"data": {
|
||||
"wsm_stall": 1262460507.897436,
|
||||
"vls_pipeline_stall": 150175218.87179488,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717094852/ISU_stall_cycles_layout20260717094940634.png"
|
||||
}
|
||||
##############################
|
||||
Sub-module: Memory Statistics
|
||||
------------------------------
|
||||
Name: VL1 Hit Rate
|
||||
Description: hit rate of all instructions in all VL1s
|
||||
Value: 87.31%
|
||||
------------------------------
|
||||
Name: L2C Hit Rate
|
||||
Description: hit rate of all instructions in all L2Cs
|
||||
Value: 97.50%
|
||||
------------------------------
|
||||
Name: Global Memory Read bytes
|
||||
Description: bytes read from global memory
|
||||
Value: 310,784,000.0byte
|
||||
------------------------------
|
||||
Name: Global Memory Write bytes
|
||||
Description: bytes write from global memory
|
||||
Value: 134,217,728.0byte
|
||||
##############################
|
||||
Sub-module: Workgroup Memory
|
||||
------------------------------
|
||||
Name: shared memory access efficiency
|
||||
Description: Proportion of NON-CONFLICT access
|
||||
Value: 74.04%
|
||||
##############################
|
||||
Sub-module: Occupancy
|
||||
------------------------------
|
||||
Name: Achieved waves
|
||||
Description: number of achieved waves
|
||||
Value: 32,768.0
|
||||
------------------------------
|
||||
Name: Dispatched waves
|
||||
Description: number of dispatched waves
|
||||
Value: 32,768.0
|
||||
##############################
|
||||
Sub-module: GPU Throughput Statistics
|
||||
------------------------------
|
||||
Name: AP MMA Duty ratio
|
||||
Description: MMA Duty ratio relative to AP active
|
||||
Value: 20.13%
|
||||
|
|
@ -0,0 +1,20 @@
|
|||
Name,Description,Value
|
||||
Total Cycles,cycles use by kernel,"52,126.52(Kcycles)"
|
||||
AP busy Duty,average AP busy duty of total cycles,99.20%
|
||||
ISU stall cycles layout,ISU stall cycles layout,"{
|
||||
""data"": {
|
||||
""wsm_stall"": 1262460507.897436,
|
||||
""vls_pipeline_stall"": 150175218.87179488,
|
||||
""vls_wdata_stall"": 2097152.0,
|
||||
""valu_stall"": 0.0
|
||||
},
|
||||
""filename"": ""/opt/mcProfiler-ubuntu18.04/output20260717094852/ISU_stall_cycles_layout20260717094940634.png""
|
||||
}"
|
||||
VL1 Hit Rate,hit rate of all instructions in all VL1s,87.31%
|
||||
L2C Hit Rate,hit rate of all instructions in all L2Cs,97.50%
|
||||
Global Memory Read bytes,bytes read from global memory,"310,784,000.0byte"
|
||||
Global Memory Write bytes,bytes write from global memory,"134,217,728.0byte"
|
||||
shared memory access efficiency,Proportion of NON-CONFLICT access,74.04%
|
||||
Achieved waves,number of achieved waves,"32,768.0"
|
||||
Dispatched waves,number of dispatched waves,"32,768.0"
|
||||
AP MMA Duty ratio,MMA Duty ratio relative to AP active,20.13%
|
||||
|
|
|
@ -0,0 +1,77 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"name": "Total Cycles",
|
||||
"description": "cycles use by kernel",
|
||||
"value": "52,126.52(Kcycles)"
|
||||
},
|
||||
{
|
||||
"name": "AP busy Duty",
|
||||
"description": "average AP busy duty of total cycles",
|
||||
"value": "99.20%"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"name": "ISU stall cycles layout",
|
||||
"description": "ISU stall cycles layout",
|
||||
"value": {
|
||||
"data": {
|
||||
"wsm_stall": 1262460507.897436,
|
||||
"vls_pipeline_stall": 150175218.87179488,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717094852/ISU_stall_cycles_layout20260717094940634.png"
|
||||
}
|
||||
}
|
||||
],
|
||||
"Memory Statistics": [
|
||||
{
|
||||
"name": "VL1 Hit Rate",
|
||||
"description": "hit rate of all instructions in all VL1s",
|
||||
"value": "87.31%"
|
||||
},
|
||||
{
|
||||
"name": "L2C Hit Rate",
|
||||
"description": "hit rate of all instructions in all L2Cs",
|
||||
"value": "97.50%"
|
||||
},
|
||||
{
|
||||
"name": "Global Memory Read bytes",
|
||||
"description": "bytes read from global memory",
|
||||
"value": "310,784,000.0byte"
|
||||
},
|
||||
{
|
||||
"name": "Global Memory Write bytes",
|
||||
"description": "bytes write from global memory",
|
||||
"value": "134,217,728.0byte"
|
||||
}
|
||||
],
|
||||
"Workgroup Memory": [
|
||||
{
|
||||
"name": "shared memory access efficiency",
|
||||
"description": "Proportion of NON-CONFLICT access",
|
||||
"value": "74.04%"
|
||||
}
|
||||
],
|
||||
"Occupancy": [
|
||||
{
|
||||
"name": "Achieved waves",
|
||||
"description": "number of achieved waves",
|
||||
"value": "32,768.0"
|
||||
},
|
||||
{
|
||||
"name": "Dispatched waves",
|
||||
"description": "number of dispatched waves",
|
||||
"value": "32,768.0"
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"name": "AP MMA Duty ratio",
|
||||
"description": "MMA Duty ratio relative to AP active",
|
||||
"value": "20.13%"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,85 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"data": 52126.522,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Total Cycles"
|
||||
},
|
||||
{
|
||||
"data": 99.19535586893751,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP busy Duty"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"data": {
|
||||
"wsm_stall": 1262460507.897436,
|
||||
"vls_pipeline_stall": 150175218.87179488,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "ISU stall cycles layout"
|
||||
}
|
||||
],
|
||||
"Memory Statistics": [
|
||||
{
|
||||
"data": 87.3124639353722,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "VL1 Hit Rate"
|
||||
},
|
||||
{
|
||||
"data": 97.4977129027295,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "L2C Hit Rate"
|
||||
},
|
||||
{
|
||||
"data": 310784000.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Global Memory Read bytes"
|
||||
},
|
||||
{
|
||||
"data": 134217728.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Global Memory Write bytes"
|
||||
}
|
||||
],
|
||||
"Workgroup Memory": [
|
||||
{
|
||||
"data": 74.03632462257644,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "shared memory access efficiency"
|
||||
}
|
||||
],
|
||||
"Occupancy": [
|
||||
{
|
||||
"data": 32768.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Achieved waves"
|
||||
},
|
||||
{
|
||||
"data": 32768.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Dispatched waves"
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"data": 20.128644256109645,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP MMA Duty ratio"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
##############################
|
||||
Sub-module: Summary
|
||||
------------------------------
|
||||
Name: Total Cycles
|
||||
Description: cycles use by kernel
|
||||
Value: 51,958.49(Kcycles)
|
||||
------------------------------
|
||||
Name: AP busy Duty
|
||||
Description: average AP busy duty of total cycles
|
||||
Value: 99.38%
|
||||
##############################
|
||||
Sub-module: ISU Statistics
|
||||
------------------------------
|
||||
Name: ISU stall cycles layout
|
||||
Description: ISU stall cycles layout
|
||||
Value: {
|
||||
"data": {
|
||||
"wsm_stall": 1230591385.6,
|
||||
"vls_pipeline_stall": 149673826.46153846,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717094852/ISU_stall_cycles_layout20260717094940987.png"
|
||||
}
|
||||
##############################
|
||||
Sub-module: Memory Statistics
|
||||
------------------------------
|
||||
Name: VL1 Hit Rate
|
||||
Description: hit rate of all instructions in all VL1s
|
||||
Value: 87.31%
|
||||
------------------------------
|
||||
Name: L2C Hit Rate
|
||||
Description: hit rate of all instructions in all L2Cs
|
||||
Value: 97.50%
|
||||
------------------------------
|
||||
Name: Global Memory Read bytes
|
||||
Description: bytes read from global memory
|
||||
Value: 310,771,712.0byte
|
||||
------------------------------
|
||||
Name: Global Memory Write bytes
|
||||
Description: bytes write from global memory
|
||||
Value: 134,217,728.0byte
|
||||
##############################
|
||||
Sub-module: Workgroup Memory
|
||||
------------------------------
|
||||
Name: shared memory access efficiency
|
||||
Description: Proportion of NON-CONFLICT access
|
||||
Value: 74.04%
|
||||
##############################
|
||||
Sub-module: Occupancy
|
||||
------------------------------
|
||||
Name: Achieved waves
|
||||
Description: number of achieved waves
|
||||
Value: 32,768.0
|
||||
------------------------------
|
||||
Name: Dispatched waves
|
||||
Description: number of dispatched waves
|
||||
Value: 32,768.0
|
||||
##############################
|
||||
Sub-module: GPU Throughput Statistics
|
||||
------------------------------
|
||||
Name: AP MMA Duty ratio
|
||||
Description: MMA Duty ratio relative to AP active
|
||||
Value: 20.15%
|
||||
|
|
@ -0,0 +1,20 @@
|
|||
Name,Description,Value
|
||||
Total Cycles,cycles use by kernel,"51,958.49(Kcycles)"
|
||||
AP busy Duty,average AP busy duty of total cycles,99.38%
|
||||
ISU stall cycles layout,ISU stall cycles layout,"{
|
||||
""data"": {
|
||||
""wsm_stall"": 1230591385.6,
|
||||
""vls_pipeline_stall"": 149673826.46153846,
|
||||
""vls_wdata_stall"": 2097152.0,
|
||||
""valu_stall"": 0.0
|
||||
},
|
||||
""filename"": ""/opt/mcProfiler-ubuntu18.04/output20260717094852/ISU_stall_cycles_layout20260717094940987.png""
|
||||
}"
|
||||
VL1 Hit Rate,hit rate of all instructions in all VL1s,87.31%
|
||||
L2C Hit Rate,hit rate of all instructions in all L2Cs,97.50%
|
||||
Global Memory Read bytes,bytes read from global memory,"310,771,712.0byte"
|
||||
Global Memory Write bytes,bytes write from global memory,"134,217,728.0byte"
|
||||
shared memory access efficiency,Proportion of NON-CONFLICT access,74.04%
|
||||
Achieved waves,number of achieved waves,"32,768.0"
|
||||
Dispatched waves,number of dispatched waves,"32,768.0"
|
||||
AP MMA Duty ratio,MMA Duty ratio relative to AP active,20.15%
|
||||
|
|
|
@ -0,0 +1,77 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"name": "Total Cycles",
|
||||
"description": "cycles use by kernel",
|
||||
"value": "51,958.49(Kcycles)"
|
||||
},
|
||||
{
|
||||
"name": "AP busy Duty",
|
||||
"description": "average AP busy duty of total cycles",
|
||||
"value": "99.38%"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"name": "ISU stall cycles layout",
|
||||
"description": "ISU stall cycles layout",
|
||||
"value": {
|
||||
"data": {
|
||||
"wsm_stall": 1230591385.6,
|
||||
"vls_pipeline_stall": 149673826.46153846,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717094852/ISU_stall_cycles_layout20260717094940987.png"
|
||||
}
|
||||
}
|
||||
],
|
||||
"Memory Statistics": [
|
||||
{
|
||||
"name": "VL1 Hit Rate",
|
||||
"description": "hit rate of all instructions in all VL1s",
|
||||
"value": "87.31%"
|
||||
},
|
||||
{
|
||||
"name": "L2C Hit Rate",
|
||||
"description": "hit rate of all instructions in all L2Cs",
|
||||
"value": "97.50%"
|
||||
},
|
||||
{
|
||||
"name": "Global Memory Read bytes",
|
||||
"description": "bytes read from global memory",
|
||||
"value": "310,771,712.0byte"
|
||||
},
|
||||
{
|
||||
"name": "Global Memory Write bytes",
|
||||
"description": "bytes write from global memory",
|
||||
"value": "134,217,728.0byte"
|
||||
}
|
||||
],
|
||||
"Workgroup Memory": [
|
||||
{
|
||||
"name": "shared memory access efficiency",
|
||||
"description": "Proportion of NON-CONFLICT access",
|
||||
"value": "74.04%"
|
||||
}
|
||||
],
|
||||
"Occupancy": [
|
||||
{
|
||||
"name": "Achieved waves",
|
||||
"description": "number of achieved waves",
|
||||
"value": "32,768.0"
|
||||
},
|
||||
{
|
||||
"name": "Dispatched waves",
|
||||
"description": "number of dispatched waves",
|
||||
"value": "32,768.0"
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"name": "AP MMA Duty ratio",
|
||||
"description": "MMA Duty ratio relative to AP active",
|
||||
"value": "20.15%"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,85 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"data": 51958.493,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Total Cycles"
|
||||
},
|
||||
{
|
||||
"data": 99.38365225488738,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP busy Duty"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"data": {
|
||||
"wsm_stall": 1230591385.6,
|
||||
"vls_pipeline_stall": 149673826.46153846,
|
||||
"vls_wdata_stall": 2097152.0,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "ISU stall cycles layout"
|
||||
}
|
||||
],
|
||||
"Memory Statistics": [
|
||||
{
|
||||
"data": 87.3124639353722,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "VL1 Hit Rate"
|
||||
},
|
||||
{
|
||||
"data": 97.49779064403101,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "L2C Hit Rate"
|
||||
},
|
||||
{
|
||||
"data": 310771712.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Global Memory Read bytes"
|
||||
},
|
||||
{
|
||||
"data": 134217728.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Global Memory Write bytes"
|
||||
}
|
||||
],
|
||||
"Workgroup Memory": [
|
||||
{
|
||||
"data": 74.03632462257644,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "shared memory access efficiency"
|
||||
}
|
||||
],
|
||||
"Occupancy": [
|
||||
{
|
||||
"data": 32768.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Achieved waves"
|
||||
},
|
||||
{
|
||||
"data": 32768.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Dispatched waves"
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"data": 20.147546446579675,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP MMA Duty ratio"
|
||||
}
|
||||
]
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 39 KiB |
|
|
@ -0,0 +1,40 @@
|
|||
{
|
||||
"RM": [
|
||||
"RM_PERF_DISPATCHED_WAVES_P3",
|
||||
"RM_PERF_DISPATCHED_WAVES_P0",
|
||||
"RM_PERF_DISPATCHED_WAVES_P1",
|
||||
"RM_PERF_DISPATCHED_WAVES_P2"
|
||||
],
|
||||
"ISU": [
|
||||
"AP_PERF_BSM_ATOMIC_CYCLES",
|
||||
"AP_PERF_ISU_ARB_VLS_STALL",
|
||||
"AP_PERF_BSM_CONFLICT_CYCLES",
|
||||
"AP_PERF_ISU_ARB_DATA_STALL",
|
||||
"AP_PERF_BSM_RD_CYCLES",
|
||||
"AP_PERF_BSM_WR_CYCLES",
|
||||
"AP_PERF_ISU_ARB_WSM_STALL",
|
||||
"AP_PERF_WAVES",
|
||||
"AP_PERF_INST_CYCLES_MMA",
|
||||
"AP_PERF_ISU_ARB_VALU_STALL",
|
||||
"AP_PERF_ISU_AP_BUSY"
|
||||
],
|
||||
"L2C": [
|
||||
"perf_dnoc_wrreq_128B",
|
||||
"perf_dnoc_wrreq_32B_low",
|
||||
"perf_dnoc_wrreq_32B_high",
|
||||
"perf_dnoc_rdreq_32B",
|
||||
"perf_hit",
|
||||
"perf_dnoc_rdreq_128B",
|
||||
"perf_dnoc_wrreq_64B",
|
||||
"perf_miss",
|
||||
"perf_dnoc_rdreq_64B",
|
||||
"perf_sector_hit"
|
||||
],
|
||||
"VL1": [
|
||||
"p_perf_thread_req_cnt",
|
||||
"p_perf_miss_cnt"
|
||||
],
|
||||
"CE": [
|
||||
"CE_PERF_BUSY_CYCLES"
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,208 @@
|
|||
{
|
||||
"perfListPerBlock": {
|
||||
"RM": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
52,
|
||||
"RM_PERF_DISPATCHED_WAVES_P3"
|
||||
],
|
||||
[
|
||||
4,
|
||||
"RM_PERF_DISPATCHED_WAVES_P0"
|
||||
],
|
||||
[
|
||||
20,
|
||||
"RM_PERF_DISPATCHED_WAVES_P1"
|
||||
],
|
||||
[
|
||||
36,
|
||||
"RM_PERF_DISPATCHED_WAVES_P2"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"ISU": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
1,
|
||||
"AP_PERF_WAVES"
|
||||
],
|
||||
[
|
||||
74,
|
||||
"AP_PERF_BSM_ATOMIC_CYCLES"
|
||||
],
|
||||
[
|
||||
80,
|
||||
"AP_PERF_ISU_ARB_VLS_STALL"
|
||||
],
|
||||
[
|
||||
78,
|
||||
"AP_PERF_BSM_CONFLICT_CYCLES"
|
||||
],
|
||||
[
|
||||
81,
|
||||
"AP_PERF_ISU_ARB_DATA_STALL"
|
||||
],
|
||||
[
|
||||
72,
|
||||
"AP_PERF_BSM_RD_CYCLES"
|
||||
],
|
||||
[
|
||||
73,
|
||||
"AP_PERF_BSM_WR_CYCLES"
|
||||
]
|
||||
]
|
||||
},
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 1,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
1,
|
||||
"AP_PERF_WAVES"
|
||||
],
|
||||
[
|
||||
79,
|
||||
"AP_PERF_ISU_ARB_WSM_STALL"
|
||||
],
|
||||
[
|
||||
43,
|
||||
"AP_PERF_INST_CYCLES_MMA"
|
||||
],
|
||||
[
|
||||
82,
|
||||
"AP_PERF_ISU_ARB_VALU_STALL"
|
||||
],
|
||||
[
|
||||
89,
|
||||
"AP_PERF_ISU_AP_BUSY"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"L2C": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
272,
|
||||
"perf_dnoc_wrreq_128B"
|
||||
],
|
||||
[
|
||||
269,
|
||||
"perf_dnoc_wrreq_32B_low"
|
||||
],
|
||||
[
|
||||
270,
|
||||
"perf_dnoc_wrreq_32B_high"
|
||||
],
|
||||
[
|
||||
46,
|
||||
"perf_dnoc_rdreq_32B"
|
||||
]
|
||||
]
|
||||
},
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 1,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
19,
|
||||
"perf_hit"
|
||||
],
|
||||
[
|
||||
48,
|
||||
"perf_dnoc_rdreq_128B"
|
||||
],
|
||||
[
|
||||
271,
|
||||
"perf_dnoc_wrreq_64B"
|
||||
],
|
||||
[
|
||||
21,
|
||||
"perf_miss"
|
||||
]
|
||||
]
|
||||
},
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 2,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
47,
|
||||
"perf_dnoc_rdreq_64B"
|
||||
],
|
||||
[
|
||||
20,
|
||||
"perf_sector_hit"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"VL1": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
79,
|
||||
"p_perf_thread_req_cnt"
|
||||
],
|
||||
[
|
||||
78,
|
||||
"p_perf_miss_cnt"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"CE": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
4,
|
||||
"CE_PERF_SEND_RM_WAVES"
|
||||
],
|
||||
[
|
||||
2,
|
||||
"CE_PERF_BUSY_CYCLES"
|
||||
]
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"action": "config,disable,read_dump,enable,kernel_collect",
|
||||
"conditions": {
|
||||
"kernelnames": [
|
||||
"packed_kernel_dense_softmax_cleanup_v12_kernel"
|
||||
],
|
||||
"include": true,
|
||||
"is_perkernel": true,
|
||||
"is_multi_device": false,
|
||||
"counts": 4
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,208 @@
|
|||
{
|
||||
"perfListPerBlock": {
|
||||
"RM": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
52,
|
||||
"RM_PERF_DISPATCHED_WAVES_P3"
|
||||
],
|
||||
[
|
||||
4,
|
||||
"RM_PERF_DISPATCHED_WAVES_P0"
|
||||
],
|
||||
[
|
||||
20,
|
||||
"RM_PERF_DISPATCHED_WAVES_P1"
|
||||
],
|
||||
[
|
||||
36,
|
||||
"RM_PERF_DISPATCHED_WAVES_P2"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"ISU": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
1,
|
||||
"AP_PERF_WAVES"
|
||||
],
|
||||
[
|
||||
74,
|
||||
"AP_PERF_BSM_ATOMIC_CYCLES"
|
||||
],
|
||||
[
|
||||
80,
|
||||
"AP_PERF_ISU_ARB_VLS_STALL"
|
||||
],
|
||||
[
|
||||
78,
|
||||
"AP_PERF_BSM_CONFLICT_CYCLES"
|
||||
],
|
||||
[
|
||||
81,
|
||||
"AP_PERF_ISU_ARB_DATA_STALL"
|
||||
],
|
||||
[
|
||||
72,
|
||||
"AP_PERF_BSM_RD_CYCLES"
|
||||
],
|
||||
[
|
||||
73,
|
||||
"AP_PERF_BSM_WR_CYCLES"
|
||||
]
|
||||
]
|
||||
},
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 1,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
1,
|
||||
"AP_PERF_WAVES"
|
||||
],
|
||||
[
|
||||
79,
|
||||
"AP_PERF_ISU_ARB_WSM_STALL"
|
||||
],
|
||||
[
|
||||
43,
|
||||
"AP_PERF_INST_CYCLES_MMA"
|
||||
],
|
||||
[
|
||||
82,
|
||||
"AP_PERF_ISU_ARB_VALU_STALL"
|
||||
],
|
||||
[
|
||||
89,
|
||||
"AP_PERF_ISU_AP_BUSY"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"L2C": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
272,
|
||||
"perf_dnoc_wrreq_128B"
|
||||
],
|
||||
[
|
||||
269,
|
||||
"perf_dnoc_wrreq_32B_low"
|
||||
],
|
||||
[
|
||||
270,
|
||||
"perf_dnoc_wrreq_32B_high"
|
||||
],
|
||||
[
|
||||
46,
|
||||
"perf_dnoc_rdreq_32B"
|
||||
]
|
||||
]
|
||||
},
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 1,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
19,
|
||||
"perf_hit"
|
||||
],
|
||||
[
|
||||
48,
|
||||
"perf_dnoc_rdreq_128B"
|
||||
],
|
||||
[
|
||||
271,
|
||||
"perf_dnoc_wrreq_64B"
|
||||
],
|
||||
[
|
||||
21,
|
||||
"perf_miss"
|
||||
]
|
||||
]
|
||||
},
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 2,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
47,
|
||||
"perf_dnoc_rdreq_64B"
|
||||
],
|
||||
[
|
||||
20,
|
||||
"perf_sector_hit"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"VL1": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
79,
|
||||
"p_perf_thread_req_cnt"
|
||||
],
|
||||
[
|
||||
78,
|
||||
"p_perf_miss_cnt"
|
||||
]
|
||||
]
|
||||
}
|
||||
],
|
||||
"CE": [
|
||||
{
|
||||
"dpcId": 0,
|
||||
"blockId": 0,
|
||||
"dpc_broadcast": false,
|
||||
"block_broadcast": false,
|
||||
"events": [
|
||||
[
|
||||
4,
|
||||
"CE_PERF_SEND_RM_WAVES"
|
||||
],
|
||||
[
|
||||
2,
|
||||
"CE_PERF_BUSY_CYCLES"
|
||||
]
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"action": "config,disable,read_dump,enable,kernel_collect",
|
||||
"conditions": {
|
||||
"kernelnames": [
|
||||
"packed_kernel_dense_softmax_cleanup_v12_kernel"
|
||||
],
|
||||
"include": true,
|
||||
"is_perkernel": true,
|
||||
"is_multi_device": false,
|
||||
"counts": 4
|
||||
}
|
||||
}
|
||||
|
|
@ -0,0 +1,249 @@
|
|||
{
|
||||
"perf_counters": {
|
||||
"0": [
|
||||
{
|
||||
"CE": {
|
||||
"generic": {
|
||||
"CE_PERF_BUSY_CYCLES": 10581470542,
|
||||
"CE_PERF_SEND_RM_WAVES": 1554177
|
||||
}
|
||||
},
|
||||
"ISU": {
|
||||
"dpcId0,blockId0": {
|
||||
"AP_PERF_WAVES": 14980,
|
||||
"AP_PERF_BSM_RD_CYCLES": 735812472,
|
||||
"AP_PERF_BSM_WR_CYCLES": 40632064,
|
||||
"AP_PERF_BSM_ATOMIC_CYCLES": 488,
|
||||
"AP_PERF_BSM_CONFLICT_CYCLES": 201592680,
|
||||
"AP_PERF_ISU_ARB_VLS_STALL": 44381341,
|
||||
"AP_PERF_ISU_ARB_DATA_STALL": 703072
|
||||
},
|
||||
"dpcId0,blockId1": {
|
||||
"AP_PERF_WAVES": 15036,
|
||||
"AP_PERF_INST_CYCLES_MMA": 1290223616,
|
||||
"AP_PERF_ISU_ARB_WSM_STALL": 372846164,
|
||||
"AP_PERF_ISU_ARB_VALU_STALL": 0,
|
||||
"AP_PERF_ISU_AP_BUSY": 1603140476
|
||||
}
|
||||
},
|
||||
"L2C": {
|
||||
"dpcId0,blockId0": {
|
||||
"perf_dnoc_wrreq_32B_low": 0,
|
||||
"perf_dnoc_wrreq_32B_high": 0,
|
||||
"perf_dnoc_wrreq_128B": 1138690,
|
||||
"perf_dnoc_rdreq_32B": 0
|
||||
},
|
||||
"dpcId0,blockId1": {
|
||||
"perf_hit": 129242083,
|
||||
"perf_miss": 3638974,
|
||||
"perf_dnoc_wrreq_64B": 0,
|
||||
"perf_dnoc_rdreq_128B": 2500143
|
||||
},
|
||||
"dpcId0,blockId2": {
|
||||
"perf_sector_hit": 2017533,
|
||||
"perf_dnoc_rdreq_64B": 268
|
||||
}
|
||||
},
|
||||
"RM": {
|
||||
"dpcId0,blockId0": {
|
||||
"RM_PERF_DISPATCHED_WAVES_P1": 0,
|
||||
"RM_PERF_DISPATCHED_WAVES_P2": 0,
|
||||
"RM_PERF_DISPATCHED_WAVES_P0": 194272,
|
||||
"RM_PERF_DISPATCHED_WAVES_P3": 0
|
||||
}
|
||||
},
|
||||
"VL1": {
|
||||
"dpcId0,blockId0": {
|
||||
"p_perf_miss_cnt": 42016808,
|
||||
"p_perf_thread_req_cnt": 331956744
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"CE": {
|
||||
"generic": {
|
||||
"CE_PERF_BUSY_CYCLES": 52126522,
|
||||
"CE_PERF_SEND_RM_WAVES": 32768
|
||||
}
|
||||
},
|
||||
"ISU": {
|
||||
"dpcId0,blockId0": {
|
||||
"AP_PERF_WAVES": 312,
|
||||
"AP_PERF_BSM_RD_CYCLES": 23743104,
|
||||
"AP_PERF_BSM_WR_CYCLES": 1310976,
|
||||
"AP_PERF_BSM_ATOMIC_CYCLES": 0,
|
||||
"AP_PERF_BSM_CONFLICT_CYCLES": 6504960,
|
||||
"AP_PERF_ISU_ARB_VLS_STALL": 1429891,
|
||||
"AP_PERF_ISU_ARB_DATA_STALL": 19968
|
||||
},
|
||||
"dpcId0,blockId1": {
|
||||
"AP_PERF_WAVES": 312,
|
||||
"AP_PERF_INST_CYCLES_MMA": 41631744,
|
||||
"AP_PERF_ISU_ARB_WSM_STALL": 12020498,
|
||||
"AP_PERF_ISU_ARB_VALU_STALL": 0,
|
||||
"AP_PERF_ISU_AP_BUSY": 51707089
|
||||
}
|
||||
},
|
||||
"L2C": {
|
||||
"dpcId0,blockId0": {
|
||||
"perf_dnoc_wrreq_32B_low": 0,
|
||||
"perf_dnoc_wrreq_32B_high": 0,
|
||||
"perf_dnoc_wrreq_128B": 32768,
|
||||
"perf_dnoc_rdreq_32B": 0
|
||||
},
|
||||
"dpcId0,blockId1": {
|
||||
"perf_hit": 4168025,
|
||||
"perf_miss": 108643,
|
||||
"perf_dnoc_wrreq_64B": 0,
|
||||
"perf_dnoc_rdreq_128B": 75871
|
||||
},
|
||||
"dpcId0,blockId2": {
|
||||
"perf_sector_hit": 65080,
|
||||
"perf_dnoc_rdreq_64B": 8
|
||||
}
|
||||
},
|
||||
"RM": {
|
||||
"dpcId0,blockId0": {
|
||||
"RM_PERF_DISPATCHED_WAVES_P1": 0,
|
||||
"RM_PERF_DISPATCHED_WAVES_P2": 0,
|
||||
"RM_PERF_DISPATCHED_WAVES_P0": 4096,
|
||||
"RM_PERF_DISPATCHED_WAVES_P3": 0
|
||||
}
|
||||
},
|
||||
"VL1": {
|
||||
"dpcId0,blockId0": {
|
||||
"p_perf_miss_cnt": 1350912,
|
||||
"p_perf_thread_req_cnt": 10647552
|
||||
}
|
||||
},
|
||||
"args": {
|
||||
"kernelName": [
|
||||
"packed_kernel_dense_softmax_cleanup_v12_kernel"
|
||||
],
|
||||
"tag": [
|
||||
-1
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
"CE": {
|
||||
"generic": {
|
||||
"CE_PERF_BUSY_CYCLES": 51958493,
|
||||
"CE_PERF_SEND_RM_WAVES": 32768
|
||||
}
|
||||
},
|
||||
"ISU": {
|
||||
"dpcId0,blockId0": {
|
||||
"AP_PERF_WAVES": 312,
|
||||
"AP_PERF_BSM_RD_CYCLES": 23743104,
|
||||
"AP_PERF_BSM_WR_CYCLES": 1310976,
|
||||
"AP_PERF_BSM_ATOMIC_CYCLES": 0,
|
||||
"AP_PERF_BSM_CONFLICT_CYCLES": 6504960,
|
||||
"AP_PERF_ISU_ARB_VLS_STALL": 1425117,
|
||||
"AP_PERF_ISU_ARB_DATA_STALL": 19968
|
||||
},
|
||||
"dpcId0,blockId1": {
|
||||
"AP_PERF_WAVES": 320,
|
||||
"AP_PERF_INST_CYCLES_MMA": 41615360,
|
||||
"AP_PERF_ISU_ARB_WSM_STALL": 12017494,
|
||||
"AP_PERF_ISU_ARB_VALU_STALL": 0,
|
||||
"AP_PERF_ISU_AP_BUSY": 51638248
|
||||
}
|
||||
},
|
||||
"L2C": {
|
||||
"dpcId0,blockId0": {
|
||||
"perf_dnoc_wrreq_32B_low": 0,
|
||||
"perf_dnoc_wrreq_32B_high": 0,
|
||||
"perf_dnoc_wrreq_128B": 32768,
|
||||
"perf_dnoc_rdreq_32B": 0
|
||||
},
|
||||
"dpcId0,blockId1": {
|
||||
"perf_hit": 4168028,
|
||||
"perf_miss": 108640,
|
||||
"perf_dnoc_wrreq_64B": 0,
|
||||
"perf_dnoc_rdreq_128B": 75868
|
||||
},
|
||||
"dpcId0,blockId2": {
|
||||
"perf_sector_hit": 65095,
|
||||
"perf_dnoc_rdreq_64B": 8
|
||||
}
|
||||
},
|
||||
"RM": {
|
||||
"dpcId0,blockId0": {
|
||||
"RM_PERF_DISPATCHED_WAVES_P1": 0,
|
||||
"RM_PERF_DISPATCHED_WAVES_P2": 0,
|
||||
"RM_PERF_DISPATCHED_WAVES_P0": 4096,
|
||||
"RM_PERF_DISPATCHED_WAVES_P3": 0
|
||||
}
|
||||
},
|
||||
"VL1": {
|
||||
"dpcId0,blockId0": {
|
||||
"p_perf_miss_cnt": 1350912,
|
||||
"p_perf_thread_req_cnt": 10647552
|
||||
}
|
||||
},
|
||||
"args": {
|
||||
"kernelName": [
|
||||
"packed_kernel_dense_softmax_cleanup_v12_kernel"
|
||||
],
|
||||
"tag": [
|
||||
-1
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"umd_data": {},
|
||||
"device_info": {
|
||||
"3": {
|
||||
"gpu_id": "51332",
|
||||
"cpu_cores_count": 0,
|
||||
"mem_banks_count": 1,
|
||||
"direct_link_count": 4,
|
||||
"indirect_link_count": 8,
|
||||
"vendor_id": 39321,
|
||||
"device_id": 16385,
|
||||
"drm_render_minor": 129,
|
||||
"isa_major": 10,
|
||||
"isa_minor": 0,
|
||||
"is_vf": 0,
|
||||
"max_vf_nums": 8,
|
||||
"domain": 0,
|
||||
"location_id": 3840,
|
||||
"num_sdma_engines": 5,
|
||||
"num_sdma_queues_per_engine": 9,
|
||||
"num_vpue_cores": 1,
|
||||
"num_vpud_cores": 8,
|
||||
"caches_count": 157,
|
||||
"peu_id_base": 0,
|
||||
"peu_count": 416,
|
||||
"max_waves_per_peu": 8,
|
||||
"wsm_size_in_kb": 64,
|
||||
"wave_front_size": 64,
|
||||
"dpc_count": 8,
|
||||
"dpc0_ap_mask": 8191,
|
||||
"dpc1_ap_mask": 73727,
|
||||
"dpc2_ap_mask": 139263,
|
||||
"dpc3_ap_mask": 204799,
|
||||
"dpc4_ap_mask": 270335,
|
||||
"dpc5_ap_mask": 335871,
|
||||
"dpc6_ap_mask": 401407,
|
||||
"dpc7_ap_mask": 466943,
|
||||
"dpc_arrays": 1,
|
||||
"ap_per_dpc": 13,
|
||||
"peu_per_ap": 4,
|
||||
"pri_mem_per_thread": 4,
|
||||
"max_slots_private_ap": 32,
|
||||
"num_ce_queues": 16,
|
||||
"max_engine_clk_gpu": 1600,
|
||||
"max_engine_clk_cpu": 3600,
|
||||
"mgpu_id": 0,
|
||||
"topology_id": 2,
|
||||
"socket_id": 3,
|
||||
"hbmecc": 1,
|
||||
"local_mem_size": 68719476736,
|
||||
"capability": 4432512,
|
||||
"maxprocess": 10
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
|
|
@ -0,0 +1,64 @@
|
|||
##############################
|
||||
Sub-module: Summary
|
||||
------------------------------
|
||||
Name: Total Cycles
|
||||
Description: cycles use by kernel
|
||||
Value: 10,581,470.54(Kcycles)
|
||||
------------------------------
|
||||
Name: AP busy Duty
|
||||
Description: average AP busy duty of total cycles
|
||||
Value: 15.15%
|
||||
##############################
|
||||
Sub-module: ISU Statistics
|
||||
------------------------------
|
||||
Name: ISU stall cycles layout
|
||||
Description: ISU stall cycles layout
|
||||
Value: {
|
||||
"data": {
|
||||
"wsm_stall": 38538769129.22506,
|
||||
"vls_pipeline_stall": 4604570054.1626835,
|
||||
"vls_wdata_stall": 72943813.86809078,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717094852/ISU_stall_cycles_layout20260717094940283.png"
|
||||
}
|
||||
##############################
|
||||
Sub-module: Memory Statistics
|
||||
------------------------------
|
||||
Name: VL1 Hit Rate
|
||||
Description: hit rate of all instructions in all VL1s
|
||||
Value: 87.34%
|
||||
------------------------------
|
||||
Name: L2C Hit Rate
|
||||
Description: hit rate of all instructions in all L2Cs
|
||||
Value: 97.30%
|
||||
------------------------------
|
||||
Name: Global Memory Read bytes
|
||||
Description: bytes read from global memory
|
||||
Value: 10,241,134,592.0byte
|
||||
------------------------------
|
||||
Name: Global Memory Write bytes
|
||||
Description: bytes write from global memory
|
||||
Value: 4,664,074,240.0byte
|
||||
##############################
|
||||
Sub-module: Workgroup Memory
|
||||
------------------------------
|
||||
Name: shared memory access efficiency
|
||||
Description: Proportion of NON-CONFLICT access
|
||||
Value: 74.04%
|
||||
##############################
|
||||
Sub-module: Occupancy
|
||||
------------------------------
|
||||
Name: Achieved waves
|
||||
Description: number of achieved waves
|
||||
Value: 1,554,177.00
|
||||
------------------------------
|
||||
Name: Dispatched waves
|
||||
Description: number of dispatched waves
|
||||
Value: 1,554,176.0
|
||||
##############################
|
||||
Sub-module: GPU Throughput Statistics
|
||||
------------------------------
|
||||
Name: AP MMA Duty ratio
|
||||
Description: MMA Duty ratio relative to AP active
|
||||
Value: 20.12%
|
||||
|
|
@ -0,0 +1,20 @@
|
|||
Name,Description,Value
|
||||
Total Cycles,cycles use by kernel,"10,581,470.54(Kcycles)"
|
||||
AP busy Duty,average AP busy duty of total cycles,15.15%
|
||||
ISU stall cycles layout,ISU stall cycles layout,"{
|
||||
""data"": {
|
||||
""wsm_stall"": 38538769129.22506,
|
||||
""vls_pipeline_stall"": 4604570054.1626835,
|
||||
""vls_wdata_stall"": 72943813.86809078,
|
||||
""valu_stall"": 0.0
|
||||
},
|
||||
""filename"": ""/opt/mcProfiler-ubuntu18.04/output20260717094852/ISU_stall_cycles_layout20260717094940283.png""
|
||||
}"
|
||||
VL1 Hit Rate,hit rate of all instructions in all VL1s,87.34%
|
||||
L2C Hit Rate,hit rate of all instructions in all L2Cs,97.30%
|
||||
Global Memory Read bytes,bytes read from global memory,"10,241,134,592.0byte"
|
||||
Global Memory Write bytes,bytes write from global memory,"4,664,074,240.0byte"
|
||||
shared memory access efficiency,Proportion of NON-CONFLICT access,74.04%
|
||||
Achieved waves,number of achieved waves,"1,554,177.00"
|
||||
Dispatched waves,number of dispatched waves,"1,554,176.0"
|
||||
AP MMA Duty ratio,MMA Duty ratio relative to AP active,20.12%
|
||||
|
|
|
@ -0,0 +1,77 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"name": "Total Cycles",
|
||||
"description": "cycles use by kernel",
|
||||
"value": "10,581,470.54(Kcycles)"
|
||||
},
|
||||
{
|
||||
"name": "AP busy Duty",
|
||||
"description": "average AP busy duty of total cycles",
|
||||
"value": "15.15%"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"name": "ISU stall cycles layout",
|
||||
"description": "ISU stall cycles layout",
|
||||
"value": {
|
||||
"data": {
|
||||
"wsm_stall": 38538769129.22506,
|
||||
"vls_pipeline_stall": 4604570054.1626835,
|
||||
"vls_wdata_stall": 72943813.86809078,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"filename": "/opt/mcProfiler-ubuntu18.04/output20260717094852/ISU_stall_cycles_layout20260717094940283.png"
|
||||
}
|
||||
}
|
||||
],
|
||||
"Memory Statistics": [
|
||||
{
|
||||
"name": "VL1 Hit Rate",
|
||||
"description": "hit rate of all instructions in all VL1s",
|
||||
"value": "87.34%"
|
||||
},
|
||||
{
|
||||
"name": "L2C Hit Rate",
|
||||
"description": "hit rate of all instructions in all L2Cs",
|
||||
"value": "97.30%"
|
||||
},
|
||||
{
|
||||
"name": "Global Memory Read bytes",
|
||||
"description": "bytes read from global memory",
|
||||
"value": "10,241,134,592.0byte"
|
||||
},
|
||||
{
|
||||
"name": "Global Memory Write bytes",
|
||||
"description": "bytes write from global memory",
|
||||
"value": "4,664,074,240.0byte"
|
||||
}
|
||||
],
|
||||
"Workgroup Memory": [
|
||||
{
|
||||
"name": "shared memory access efficiency",
|
||||
"description": "Proportion of NON-CONFLICT access",
|
||||
"value": "74.04%"
|
||||
}
|
||||
],
|
||||
"Occupancy": [
|
||||
{
|
||||
"name": "Achieved waves",
|
||||
"description": "number of achieved waves",
|
||||
"value": "1,554,177.00"
|
||||
},
|
||||
{
|
||||
"name": "Dispatched waves",
|
||||
"description": "number of dispatched waves",
|
||||
"value": "1,554,176.0"
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"name": "AP MMA Duty ratio",
|
||||
"description": "MMA Duty ratio relative to AP active",
|
||||
"value": "20.12%"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,85 @@
|
|||
{
|
||||
"Summary": [
|
||||
{
|
||||
"data": 10581470.542,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Total Cycles"
|
||||
},
|
||||
{
|
||||
"data": 15.150450682982205,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP busy Duty"
|
||||
}
|
||||
],
|
||||
"ISU Statistics": [
|
||||
{
|
||||
"data": {
|
||||
"wsm_stall": 38538769129.22506,
|
||||
"vls_pipeline_stall": 4604570054.1626835,
|
||||
"vls_wdata_stall": 72943813.86809078,
|
||||
"valu_stall": 0.0
|
||||
},
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "ISU stall cycles layout"
|
||||
}
|
||||
],
|
||||
"Memory Statistics": [
|
||||
{
|
||||
"data": 87.34268582897053,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "VL1 Hit Rate"
|
||||
},
|
||||
{
|
||||
"data": 97.30243733459334,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "L2C Hit Rate"
|
||||
},
|
||||
{
|
||||
"data": 10241134592.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Global Memory Read bytes"
|
||||
},
|
||||
{
|
||||
"data": 4664074240.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Global Memory Write bytes"
|
||||
}
|
||||
],
|
||||
"Workgroup Memory": [
|
||||
{
|
||||
"data": 74.03645154920846,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "shared memory access efficiency"
|
||||
}
|
||||
],
|
||||
"Occupancy": [
|
||||
{
|
||||
"data": 1554176.9999999998,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Achieved waves"
|
||||
},
|
||||
{
|
||||
"data": 1554176.0,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "Dispatched waves"
|
||||
}
|
||||
],
|
||||
"GPU Throughput Statistics": [
|
||||
{
|
||||
"data": 20.12025201963649,
|
||||
"isError": false,
|
||||
"message": "",
|
||||
"name": "AP MMA Duty ratio"
|
||||
}
|
||||
]
|
||||
}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,16 @@
|
|||
case_id,config,batch,total_q,total_kv,max_q,max_kv,seq_len,compile_first_launch_s,pilot_ms,repeats_per_sample,sample_count,median_ms,p10_ms,p90_ms,min_ms,max_ms,spread_pct
|
||||
1,ragged_b33_total16294,33,16294,16294,987,987,987,4.095663,2.072576,48,15,2.050485,2.049415,2.051076,2.048971,2.051813,0.081
|
||||
2,equal_b1_s1024,1,1024,1024,1024,1024,1024,0.209597,0.309931,322,15,0.290686,0.290553,0.290940,0.290066,0.291083,0.133
|
||||
3,equal_b1_s4096,1,4096,4096,4096,4096,4096,0.210687,3.035733,32,15,3.018648,3.017800,3.019238,3.017704,3.019256,0.048
|
||||
4,equal_b1_s16384,1,16384,16384,16384,16384,16384,0.251069,45.958740,2,15,45.972607,45.957453,45.993779,45.947903,46.003456,0.079
|
||||
5,equal_b4_s1024,4,4096,4096,1024,1024,1024,0.204058,0.921856,108,15,0.906871,0.905927,0.907311,0.905169,0.907750,0.153
|
||||
6,equal_b4_s4096,4,16384,16384,4096,4096,4096,0.222118,12.360874,8,15,12.336992,12.333901,12.340026,12.332096,12.340640,0.050
|
||||
7,equal_b16_s1024,16,16384,16384,1024,1024,1024,0.206309,3.361109,29,15,3.348056,3.347334,3.350782,3.346600,3.352947,0.103
|
||||
8,equal_b16_s2048,16,32768,32768,2048,2048,2048,0.292793,13.144918,7,15,13.120476,13.118625,13.123906,13.117549,13.132983,0.040
|
||||
9,varlen_uniform_q512_k1024_b4,4,2048,4096,512,1024,1024,0.215894,0.737451,135,15,0.717845,0.717765,0.717965,0.717282,0.718033,0.028
|
||||
10,varlen_mixed_b4,4,1536,3584,640,1280,1280,0.212858,0.552704,180,15,0.533872,0.533763,0.534045,0.533238,0.534074,0.053
|
||||
11,varlen_q_lt_kv_b2,2,1024,3072,512,2048,2048,0.215189,0.603221,165,15,0.584819,0.584664,0.585137,0.584329,0.585174,0.081
|
||||
12,ragged_b27_total12251,27,12251,12251,873,873,873,0.216958,1.495125,66,15,1.478202,1.477512,1.478641,1.477295,1.478931,0.076
|
||||
13,short_ragged_b15_total969,15,969,969,123,123,123,0.176575,0.288768,346,15,0.272169,0.272088,0.272253,0.271831,0.272271,0.061
|
||||
14,single_token,1,1,1,1,1,1,0.170628,0.041813,2000,15,0.030645,0.030608,0.031089,0.030598,0.031225,1.567
|
||||
15,tail_non_power2,2,98,98,65,65,65,0.169383,0.069803,1432,15,0.051921,0.051709,0.051938,0.051699,0.051939,0.441
|
||||
|
|
|
@ -0,0 +1,407 @@
|
|||
* Serving Flask app 'Profiler' (lazy loading)
|
||||
* Environment: production
|
||||
WARNING: This is a development server. Do not use it in a production deployment.
|
||||
Use a production WSGI server instead.
|
||||
* Debug mode: off
|
||||
profiler http listen to 127.0.0.1:50123
|
||||
try to connect to server...
|
||||
connect to server success
|
||||
cmdline is: python tests/trace_tilelang_64g_case4.py
|
||||
casename is: opt012_case4_targeted
|
||||
metrics is: ['Total Cycles', 'AP busy Duty', 'AP MMA Duty ratio', 'ISU stall cycles layout', 'VL1 Hit Rate', 'L2C Hit Rate', 'Global Memory Read bytes', 'Global Memory Write bytes', 'shared memory access efficiency', 'Achieved waves', 'Dispatched waves']
|
||||
start a local task...
|
||||
[info] start a new perf exec thread
|
||||
[info] start perf_exec
|
||||
[info] new exec 054b0685-1efd-4049-82d2-21af8938b66e:('python tests/trace_tilelang_64g_case4.py', 'opt012_case4_targeted', ['Total Cycles', 'AP busy Duty', 'AP MMA Duty ratio', 'ISU stall cycles layout', 'VL1 Hit Rate', 'L2C Hit Rate', 'Global Memory Read bytes', 'Global Memory Write bytes', 'shared memory access efficiency', 'Achieved waves', 'Dispatched waves'])
|
||||
########## exec_id: 054b0685-1efd-4049-82d2-21af8938b66e ##########
|
||||
[info] start generate batch files
|
||||
[info] output path is: /opt/mcProfiler-ubuntu18.04/output20260717094852
|
||||
[info] generate event_batch_0.json successfully!
|
||||
[info] generate batch files done
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
Loading tilelang libs from dev root: /data/tilelang-metax/build
|
||||
[09:48:54.609][MXS][I]| MACA_LAUNCH_BLOCKING | 1 | 0 |
|
||||
[09:48:54.610][MXS][W]| MCTX_TARGET_INIT | 1 | 0 |
|
||||
[09:48:54.612][MCR][I]mc_device.cpp :562 : Signal SIGUSR1 now has a custom handler to enable mctx at runtime!
|
||||
[09:48:54.612][MCR][I]mc_device.cpp :309 : mcSnapShotCtx: Attempting to open shared memory: /shm_snapshot_100986
|
||||
mctxAutoStart called!
|
||||
[09:48:54.612][MCTX][I]mctxTPImpl.cpp :201 : MACA configuration file path redirected from ${MACA_PATH}/etc to MACA_ETC_PATH
|
||||
[09:48:54.612][MCTX][I]mctxTPImpl.cpp :201 : MACA configuration file path redirected from ${MACA_PATH}/etc to MACA_ETC_PATH
|
||||
[09:48:54.612][MCTX][I]mctxTPImpl.cpp :300 : profiler target ENV file[/opt/mcProfiler-ubuntu18.04/output20260717094852/mcProfiler.json] exist!
|
||||
[09:48:54.613][MCR][I]mc_device.cpp :849 : Init enter!
|
||||
[09:48:54.613][MXC][I]MxcInit: entering
|
||||
[09:48:54.627][MXKW][I]topology.c :1358: [topology_sysfs_get_system_props]PlatformCpuNodes:2 ,num_sysfs_nodes:10
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/2/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 1,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/4/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 2,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/5/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 3,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/6/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 4,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/7/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 5,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/8/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 6,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/9/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 7,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :1460: [topology_sysfs_get_system_props]cpu total 2 nodes, map list:
|
||||
[09:48:54.627][MXKW][I]topology.c :1463: [topology_sysfs_get_system_props]map_cpu[0]:0
|
||||
[09:48:54.627][MXKW][I]topology.c :1463: [topology_sysfs_get_system_props]map_cpu[1]:1
|
||||
[09:48:54.627][MXKW][I]topology.c :1465: [topology_sysfs_get_system_props]gpu total 1 nodes, map list:
|
||||
[09:48:54.627][MXKW][I]topology.c :1468: [topology_sysfs_get_system_props]map_gpu[0]:3
|
||||
[09:48:54.627][MXKW][I]topology.c :1516: [topology_sysfs_init_gpu_page_size]XCORE_PAGE_SIZE: 0x200000, GPU_PAGE_SIZE: 0x200000
|
||||
[09:48:54.627][MXKW][I]topology.c :954 : [find_pci_info]create pci_info_array. It has 10 element. directory_name is 0.
|
||||
[09:48:54.627][MXKW][I]topology.c :962 : [find_pci_info]node 0 's pci info is not found
|
||||
[09:48:54.627][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/0/location file!
|
||||
[09:48:54.627][MXKW][I]topology.c :962 : [find_pci_info]node 0 's pci info is not found
|
||||
[09:48:54.627][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/0/location file!
|
||||
[09:48:54.627][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 0, isa_minor is 0, MXName is
|
||||
[09:48:54.627][MXKW][I]topology.c :962 : [find_pci_info]node 1 's pci info is not found
|
||||
[09:48:54.627][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/1/location file!
|
||||
[09:48:54.627][MXKW][I]topology.c :962 : [find_pci_info]node 1 's pci info is not found
|
||||
[09:48:54.627][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/1/location file!
|
||||
[09:48:54.627][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 0, isa_minor is 0, MXName is
|
||||
[09:48:54.627][MXKW][I]topology.c :962 : [find_pci_info]node 3 's pci info is not found
|
||||
[09:48:54.627][MXKW][I]topology.c :984 : [add_pci_info]added node 3 pci info , bdf_value:3840, domain_value:0
|
||||
[09:48:54.627][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 10, isa_minor is 0, MXName is
|
||||
[09:48:54.627][MXKW][I]fmm.c :2445: [init_mmap_apertures]Initialized unreserved SVM apertures: 0xa00000000 - 0x7fffffffffff
|
||||
[09:48:54.627][MXKW][I]topology.c :1358: [topology_sysfs_get_system_props]PlatformCpuNodes:2 ,num_sysfs_nodes:10
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/2/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 1,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/4/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 2,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/5/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 3,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/6/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 4,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/7/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 5,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/8/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 6,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/9/gpu_id,res: 0,errno:1
|
||||
[09:48:54.627][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 7,ret is 0,is_node_supported is 0
|
||||
[09:48:54.627][MXKW][I]topology.c :1460: [topology_sysfs_get_system_props]cpu total 2 nodes, map list:
|
||||
[09:48:54.627][MXKW][I]topology.c :1463: [topology_sysfs_get_system_props]map_cpu[0]:0
|
||||
[09:48:54.627][MXKW][I]topology.c :1463: [topology_sysfs_get_system_props]map_cpu[1]:1
|
||||
[09:48:54.627][MXKW][I]topology.c :1465: [topology_sysfs_get_system_props]gpu total 1 nodes, map list:
|
||||
[09:48:54.627][MXKW][I]topology.c :1468: [topology_sysfs_get_system_props]map_gpu[0]:3
|
||||
[09:48:54.627][MXKW][I]topology.c :962 : [find_pci_info]node 0 's pci info is not found
|
||||
[09:48:54.627][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/0/location file!
|
||||
[09:48:54.627][MXKW][I]topology.c :962 : [find_pci_info]node 0 's pci info is not found
|
||||
[09:48:54.627][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/0/location file!
|
||||
[09:48:54.627][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 0, isa_minor is 0, MXName is
|
||||
[09:48:54.628][MXKW][I]topology.c :2548: [topology_take_snapshot]node0 direct_links route:
|
||||
[09:48:54.628][MXKW][I]topology.c :2562: [topology_take_snapshot]link0 type:2, 0--->1
|
||||
[09:48:54.628][MXKW][I]topology.c :2562: [topology_take_snapshot]link1 type:3, 0--->2
|
||||
[09:48:54.628][MXKW][I]topology.c :962 : [find_pci_info]node 1 's pci info is not found
|
||||
[09:48:54.628][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/1/location file!
|
||||
[09:48:54.628][MXKW][I]topology.c :962 : [find_pci_info]node 1 's pci info is not found
|
||||
[09:48:54.628][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/1/location file!
|
||||
[09:48:54.628][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 0, isa_minor is 0, MXName is
|
||||
[09:48:54.628][MXKW][I]topology.c :2548: [topology_take_snapshot]node1 direct_links route:
|
||||
[09:48:54.628][MXKW][I]topology.c :2562: [topology_take_snapshot]link0 type:2, 1--->0
|
||||
[09:48:54.628][MXKW][I]topology.c :2573: [topology_take_snapshot]node1 indirect_links route:
|
||||
[09:48:54.628][MXKW][I]topology.c :2597: [topology_take_snapshot]link1 type:105 1--->2
|
||||
[09:48:54.628][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 10, isa_minor is 0, MXName is
|
||||
[09:48:54.630][MXKW][I]topology.c :2548: [topology_take_snapshot]node2 direct_links route:
|
||||
[09:48:54.630][MXKW][I]topology.c :2562: [topology_take_snapshot]link0 type:3, 2--->0
|
||||
[09:48:54.630][MXKW][I]topology.c :2573: [topology_take_snapshot]node2 indirect_links route:
|
||||
[09:48:54.630][MXKW][I]topology.c :2597: [topology_take_snapshot]link1 type:105 2--->1
|
||||
[09:48:54.630][MXC][I]MxcInit: leaving
|
||||
[09:48:54.630][MXC][I]InitDma: Entering
|
||||
[09:48:54.634][MXC][I]InitDma: Leaving
|
||||
[09:48:54.634][MCR][I]mx_device.cpp :2294: Device xcore1000 create, major:10, minor:0
|
||||
[09:48:54.634][MCR][I]mx_device.cpp :2795: private_mem:436207616, info_.globalMemSize_:68283269120, gpu_single_alloc_percent:0.990000
|
||||
[09:48:54.634][MCR][I]mx_device.cpp :2865: get registersPerMtreg_:64,info_dpcNum_:8,info_wavePerDpc_:416, wavefrontWidth_:64
|
||||
[09:48:54.634][MCR][I]mx_device.cpp :3042: rmNum_:8, isuNum_:13, sl1Num_:7, vlsNum_:13, vl1Num_:13, l2cNum_:32, l2aNum_:4, l2xNum_:4
|
||||
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :1083: getQueueCnt ce queue bitmap:ffff
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :1091: getQueueCnt hwQueueNum:16
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :1153: agent_handle=0x55f372b7d800,hardware queues number:16, queue count configuration: 8-4-4(H-N-L), highForMccl:0
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :2445: Device xcore1000 create fininshed!
|
||||
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :2294: Device xcore1000 create, major:10, minor:0
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :2795: private_mem:436207616, info_.globalMemSize_:68283269120, gpu_single_alloc_percent:0.990000
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :2865: get registersPerMtreg_:64,info_dpcNum_:8,info_wavePerDpc_:416, wavefrontWidth_:64
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :3042: rmNum_:8, isuNum_:13, sl1Num_:7, vlsNum_:13, vl1Num_:13, l2cNum_:32, l2aNum_:4, l2xNum_:4
|
||||
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :1083: getQueueCnt ce queue bitmap:ffff
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :1091: getQueueCnt hwQueueNum:16
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :1153: agent_handle=0x55f372b7d800,hardware queues number:16, queue count configuration: 8-4-4(H-N-L), highForMccl:0
|
||||
[09:48:54.635][MCR][I]mx_device.cpp :2445: Device xcore1000 create fininshed!
|
||||
|
||||
[09:48:54.636][MCR][I]mc_vpu_impl.cpp :80 : mcVpuLoad mxvpu lib version: 1.0.20260512_081136
|
||||
|
||||
[09:48:54.636][MCR][I]mc_vpu_impl.cpp :86 : mcVpuLoad load mxvpu lib path: /opt/maca-3.7.1/lib/libmxvpu.so
|
||||
|
||||
[09:48:54.636][MCC][I]AddDevice device_id:0, arch:xcore1000
|
||||
[09:48:54.636][MCC][I]Get compiler from environment variable MACA_CLANG_PATH: /opt/maca/mxgpu_llvm/bin.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[09:49:16.482][MXS][I] Node0 Node1
|
||||
[09:49:16.482][MXS][I]Gpu0 X
|
||||
[09:49:16.482][MXS][I]CpuBind X X
|
||||
get_mempolicy: Operation not permitted
|
||||
[09:49:16.482][MXS][I]MemBind X X
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[09:49:17.300][MXS][I]mcruntime 100% 0%
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[09:49:18.121][MXS][I]mccompiler 100% 0%
|
||||
[09:49:18.941][MXS][I]mxc-runtime 99% 1%
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[09:49:19.758][MXS][I]mxkw 100%
|
||||
[09:49:19.758][PTI][I]mcpti_profiler.cpp :618 : getModuleInfo:1,8,13,7,13,13,32,4,4
|
||||
|
||||
[09:49:19.760][MCR][I]mx_device.cpp :5792: agent_handle=0x55f372b7d800
|
||||
[09:49:19.760][MXKW][I]queues.c :388 : [get_quantum_map_value]clockGrain 4
|
||||
[09:49:19.760][MXKW][I]queues.c :393 : [get_quantum_map_value]highQuantum 16 normal 8 low 4
|
||||
[09:49:19.761][MCR][I]mx_device.cpp :1446: created hardware queue 0x7faf2da1e000(HWq=0x7fab3c001cd0) with size 1024 with priority 0(L), cooperative: 0, is_exclusive: 0
|
||||
[09:49:19.761][MCR][I]mx_device.cpp :1458: number of allocated hardware queues 0-0-1(H-N-L), number of users 0-0-1
|
||||
[09:49:19.764][MCR][W]mx_perf_counter.cpp :568 : Perf is not exclusive, can not be disabled!
|
||||
[09:49:19.764][MCR][W]mx_perf_counter.cpp :568 : Perf is not exclusive, can not be disabled!
|
||||
[09:49:19.765][MCTX][I]mctxTPImpl.cpp :186 : Automatically acquire a port.
|
||||
[09:49:19.765][MCTX][I]mcRpcAsyncImpl.cpp :421 : mcRpcServerAsyncInit 0x55f372b6a2b0 0
|
||||
[09:49:19.765][MCTX][I]mcRpcAsyncImpl.cpp :500 : mcRpcAsyncServerTask start!
|
||||
[09:49:19.765][MCTX][I]mcRpcAsyncImpl.cpp :472 : mcRpcServerAsyncRun start
|
||||
[09:49:19.766][MCTX][I]server.cpp :69 : Server thread num: 1, listening on: 39971.
|
||||
[info] MctxStreamProfilerCountDataGet Start
|
||||
[info] MctxStreamProfilerCountDataGet Loop
|
||||
mcToolsExtPid_UsedPort:100986-39971
|
||||
[09:49:23.465][MXC][I]raw dispatch ts:[252024761964536,252024762027490]
|
||||
[09:49:23.465][MXC][I]dispatch ts translate ret [1784281763465289495, 1784281763465919033]
|
||||
[09:49:23.466][MXC][I]raw dispatch ts:[252024762046270,252024762055562]
|
||||
[09:49:23.466][MXC][I]dispatch ts translate ret [1784281763466106883, 1784281763466199838]
|
||||
[09:49:23.466][MXC][I]raw dispatch ts:[252024762075727,252024762084970]
|
||||
[09:49:23.466][MXC][I]dispatch ts translate ret [1784281763466401432, 1784281763466493831]
|
||||
[09:49:23.466][MCR][I]mc_memory.cpp :1671: copy strategy - H2DCpuCopy:0,D2HCpuCopy:0,cpuCopy:0,d2dCrossDevices:0,pageableHToDAsync:0
|
||||
[09:49:23.466][MCR][I]mc_memory.cpp :1672: copy strategy - isAsync:0,waitComplete:1
|
||||
[09:49:23.466][MXC][I]GpuDevice 0x72b7d800 enable profiling, sdm_counter:202345146698043, sys_time:1784281763466835437
|
||||
[09:49:23.466][MXC][I]raw async_copy ts:[202345146719485,202345146722136]
|
||||
[09:49:23.466][MXC][I]sdma_ts_.sdma_counter:202345146698043
|
||||
[09:49:23.466][MXC][I]DNOCfreq:900000000, async_copy ts translate ret [1784281763466859261, 1784281763466862206]
|
||||
[info] MctxStreamProfilerCountDataGet Loop
|
||||
[09:49:27.561][MXC][I]raw dispatch ts:[252025166946751,252025171544787]
|
||||
[09:49:27.561][MXC][I]dispatch ts translate ret [1784281767515100424, 1784281767561080657]
|
||||
[09:49:27.607][MXC][I]raw dispatch ts:[252025171612630,252025176205174]
|
||||
[09:49:27.607][MXC][I]dispatch ts translate ret [1784281767561759085, 1784281767607684381]
|
||||
[09:49:27.653][MXC][I]raw dispatch ts:[252025176237850,252025180832045]
|
||||
[09:49:27.653][MXC][I]dispatch ts translate ret [1784281767608011140, 1784281767653952975]
|
||||
[09:49:27.700][MXC][I]raw dispatch ts:[252025180847484,252025185443434]
|
||||
[09:49:27.700][MXC][I]dispatch ts translate ret [1784281767654107364, 1784281767700066737]
|
||||
[info] MctxStreamProfilerCountDataGet Loop
|
||||
[09:49:27.746][MXC][I]raw dispatch ts:[252025185454245,252025190049041]
|
||||
[09:49:27.746][MXC][I]dispatch ts translate ret [1784281767700174846, 1784281767746122668]
|
||||
[09:49:27.792][MXC][I]raw dispatch ts:[252025190059345,252025194653531]
|
||||
[09:49:27.792][MXC][I]dispatch ts translate ret [1784281767746225707, 1784281767792167451]
|
||||
[09:49:27.838][MXC][I]raw dispatch ts:[252025194664186,252025199259399]
|
||||
[09:49:27.838][MXC][I]dispatch ts translate ret [1784281767792274000, 1784281767838225993]
|
||||
[09:49:27.884][MXC][I]raw dispatch ts:[252025199269658,252025203864078]
|
||||
[09:49:27.884][MXC][I]dispatch ts translate ret [1784281767838328582, 1784281767884272657]
|
||||
[09:49:27.930][MXC][I]raw dispatch ts:[252025203875030,252025208469093]
|
||||
[09:49:27.930][MXC][I]dispatch ts translate ret [1784281767884382176, 1784281767930322678]
|
||||
[09:49:27.976][MXC][I]raw dispatch ts:[252025208479330,252025213073049]
|
||||
[09:49:27.976][MXC][I]dispatch ts translate ret [1784281767930425047, 1784281767976362108]
|
||||
[09:49:28.022][MXC][I]raw dispatch ts:[252025213083311,252025217678740]
|
||||
[09:49:28.022][MXC][I]dispatch ts translate ret [1784281767976464728, 1784281768022419179]
|
||||
TRACE_REGION_BEGIN repeats=20
|
||||
[09:49:28.068][MXC][I]raw dispatch ts:[252025217695726,252025222288192]
|
||||
[09:49:28.068][MXC][I]dispatch ts translate ret [1784281768022589038, 1784281768068513296]
|
||||
[09:49:28.114][MXC][I]raw dispatch ts:[252025222301060,252025226895822]
|
||||
[09:49:28.114][MXC][I]dispatch ts translate ret [1784281768068641975, 1784281768114589469]
|
||||
[09:49:28.160][MXC][I]raw dispatch ts:[252025226910125,252025231501698]
|
||||
[09:49:28.160][MXC][I]dispatch ts translate ret [1784281768114732498, 1784281768160648098]
|
||||
[09:49:28.206][MXC][I]raw dispatch ts:[252025231511388,252025236106131]
|
||||
[09:49:28.206][MXC][I]dispatch ts translate ret [1784281768160744997, 1784281768206692290]
|
||||
[09:49:28.252][MXC][I]raw dispatch ts:[252025236116876,252025240710679]
|
||||
[09:49:28.252][MXC][I]dispatch ts translate ret [1784281768206799739, 1784281768252737652]
|
||||
[09:49:28.298][MXC][I]raw dispatch ts:[252025240720750,252025245316326]
|
||||
[09:49:28.298][MXC][I]dispatch ts translate ret [1784281768252838361, 1784281768298793984]
|
||||
[09:49:28.344][MXC][I]raw dispatch ts:[252025245326523,252025249922572]
|
||||
[09:49:28.344][MXC][I]dispatch ts translate ret [1784281768298895953, 1784281768344856325]
|
||||
[09:49:28.390][MXC][I]raw dispatch ts:[252025249932766,252025254524670]
|
||||
[09:49:28.390][MXC][I]dispatch ts translate ret [1784281768344958264, 1784281768390877177]
|
||||
[09:49:28.436][MXC][I]raw dispatch ts:[252025254534745,252025259130791]
|
||||
[09:49:28.436][MXC][I]dispatch ts translate ret [1784281768390977926, 1784281768436938259]
|
||||
[09:49:28.483][MXC][I]raw dispatch ts:[252025259142769,252025263736648]
|
||||
[09:49:28.483][MXC][I]dispatch ts translate ret [1784281768437058038, 1784281768482996703]
|
||||
[09:49:28.529][MXC][I]raw dispatch ts:[252025263748950,252025268343798]
|
||||
[09:49:28.529][MXC][I]dispatch ts translate ret [1784281768483119722, 1784281768529068064]
|
||||
[09:49:28.575][MXC][I]raw dispatch ts:[252025268354282,252025272947734]
|
||||
[09:49:28.575][MXC][I]dispatch ts translate ret [1784281768529172903, 1784281768575107306]
|
||||
[09:49:28.621][MXC][I]raw dispatch ts:[252025272958606,252025277553670]
|
||||
[09:49:28.621][MXC][I]dispatch ts translate ret [1784281768575216025, 1784281768621166537]
|
||||
[09:49:28.667][MXC][I]raw dispatch ts:[252025277563357,252025282158967]
|
||||
[09:49:28.667][MXC][I]dispatch ts translate ret [1784281768621263406, 1784281768667219389]
|
||||
[09:49:28.713][MXC][I]raw dispatch ts:[252025282168721,252025286763403]
|
||||
[09:49:28.713][MXC][I]dispatch ts translate ret [1784281768667316928, 1784281768713263611]
|
||||
[09:49:28.759][MXC][I]raw dispatch ts:[252025286774296,252025291371592]
|
||||
[09:49:28.759][MXC][I]dispatch ts translate ret [1784281768713372541, 1784281768759345531]/tmp/_MEI4fqWsX/phttp/http_server.py:89: UserWarning:
|
||||
|
||||
The 'environ['werkzeug.server.shutdown']' function is deprecated and will be removed in Werkzeug 2.1.
|
||||
|
||||
|
||||
[09:49:28.805][MXC][I]raw dispatch ts:[252025291381474,252025295975968]
|
||||
[09:49:28.805][MXC][I]dispatch ts translate ret [1784281768759444351, 1784281768805389014]
|
||||
[09:49:28.851][MXC][I]raw dispatch ts:[252025295986210,252025300579049]
|
||||
[09:49:28.851][MXC][I]dispatch ts translate ret [1784281768805491433, 1784281768851419688]
|
||||
[09:49:28.897][MXC][I]raw dispatch ts:[252025300590839,252025305184548]
|
||||
[09:49:28.897][MXC][I]dispatch ts translate ret [1784281768851537587, 1784281768897474550]
|
||||
[09:49:28.943][MXC][I]raw dispatch ts:[252025305196876,252025309791449]
|
||||
[09:49:28.943][MXC][I]dispatch ts translate ret [1784281768897597829, 1784281768943543442]
|
||||
TRACE_REGION_END
|
||||
[09:49:28.944][MXC][I]raw dispatch ts:[252025309881486,252025309900362]
|
||||
[09:49:28.944][MXC][I]dispatch ts translate ret [1784281768944443884, 1784281768944632659]
|
||||
[09:49:28.946][MXC][I]raw dispatch ts:[252025310035564,252025310050351]
|
||||
[09:49:28.946][MXC][I]dispatch ts translate ret [1784281768945984566, 1784281768946132424]
|
||||
[09:49:28.946][MXC][I]raw dispatch ts:[252025310070673,252025310085532]
|
||||
[09:49:28.946][MXC][I]dispatch ts translate ret [1784281768946335649, 1784281768946484244]
|
||||
[09:49:28.948][MXC][I]raw dispatch ts:[252025310248003,252025310262412]
|
||||
[09:49:28.948][MXC][I]dispatch ts translate ret [1784281768948108944, 1784281768948253034]
|
||||
[09:49:28.948][MXC][I]raw dispatch ts:[252025310295859,252025310296366]
|
||||
[09:49:28.948][MXC][I]dispatch ts translate ret [1784281768948587500, 1784281768948592570]
|
||||
[09:49:28.950][MXC][I]raw dispatch ts:[252025310515936,252025310523599]
|
||||
[09:49:28.950][MXC][I]dispatch ts translate ret [1784281768950788250, 1784281768950864880]
|
||||
[09:49:28.951][MCR][I]mc_memory.cpp :1671: copy strategy - H2DCpuCopy:0,D2HCpuCopy:0,cpuCopy:0,d2dCrossDevices:0,pageableHToDAsync:0
|
||||
[09:49:28.951][MCR][I]mc_memory.cpp :1672: copy strategy - isAsync:0,waitComplete:0
|
||||
[09:49:28.951][MXC][I]raw async_copy ts:[202351727798658,202351727800587]
|
||||
[09:49:28.951][MXC][I]sdma_ts_.sdma_counter:202345146698043
|
||||
[09:49:28.951][MXC][I]DNOCfreq:900000000, async_copy ts translate ret [1784281770779169453, 1784281770779171596]
|
||||
[09:49:28.951][MCR][W]mx_command.cpp :1883: HW exec timestamp will be adjusted!!!
|
||||
[09:49:28.951][MCR][W]mx_command.cpp :1885: floor:1784281768951060820, start:1784281770779169536, end:1784281770779171584, ceiling:1784281768951080722
|
||||
TILELANG_64G_CASE4_TRACE_PASS
|
||||
[09:49:29.173][MCR][I]mc_device.cpp :1002: tearDown enter!
|
||||
|
||||
[09:49:29.173][MCR][I]mx_device.cpp :1578: deleting hardware queue 0x7faf2da1e000 with refCount 0
|
||||
[09:49:29.173][MCR][I]mx_device.cpp :1601: number of allocated hardware queues 0-0-0(H-N-L), number of users 0-0-0
|
||||
[09:49:29.174][MXC][I]~GpuDevice: Entering
|
||||
[09:49:29.174][MCTX][I]mcRpcAsyncImpl.cpp :550 : callback mcRpcAsyncDone!
|
||||
[09:49:29.174][PTI][I]mcpti_tracer.hpp :309 : stopProcessThread
|
||||
[09:49:29.174][MCTX][I]mcRpcAsyncImpl.cpp :491 : mcRpcServerAsync TearDown!
|
||||
[09:49:29.174][MCTX][I]server.cpp :116 : Server will tearDown[3].
|
||||
[09:49:29.181][MCTX][I]mcRpcAsyncImpl.cpp :250 : ProfilerCountDataGet the last msg:[
|
||||
{
|
||||
"CE" : [
|
||||
{
|
||||
"blockId" : 0,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"2" : 175708168690490,
|
||||
"4" : 450071379102
|
||||
}
|
||||
}
|
||||
],
|
||||
"ISU" : [
|
||||
{
|
||||
"blockId" : 0,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"1" : 1212512628467,
|
||||
"72" : 1786425550280,
|
||||
"73" : 1347731440785,
|
||||
"74" : 1246633206573,
|
||||
"78" : 7260895072502,
|
||||
"80" : 9324991417010,
|
||||
"81" : 1362363000338
|
||||
}
|
||||
},
|
||||
{
|
||||
"blockId" : 1,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"1" : 1212538968661,
|
||||
"43" : 1255515835759,
|
||||
"79" : 19719449634867,
|
||||
"82" : 16744769089836,
|
||||
"89" : 2422296694583
|
||||
}
|
||||
}
|
||||
],
|
||||
"L2C" : [
|
||||
{
|
||||
"blockId" : 0,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"269" : 24165782165,
|
||||
"270" : 10537367113,
|
||||
"272" : 365067716547,
|
||||
"46" : 734181528377
|
||||
}
|
||||
},
|
||||
{
|
||||
[info] MctxStreamProfilerCountDataGet Loop
|
||||
"blockId" : 1,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"19" : 1358922567291,
|
||||
"21" : 3611410135643,
|
||||
"271" : 776485210,
|
||||
"48" : 3342619794
|
||||
}
|
||||
},
|
||||
{
|
||||
"blockId" : 2,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"20" : 1920936811968,
|
||||
"47" : 1025093557705
|
||||
}
|
||||
}
|
||||
],
|
||||
"RM" : [
|
||||
{
|
||||
"blockId" : 0,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"20" : 0,
|
||||
"36" : 0,
|
||||
"4" : 194453959131,
|
||||
"52" : 0
|
||||
}
|
||||
}
|
||||
],
|
||||
"VL1" : [
|
||||
{
|
||||
"blockId" : 0,
|
||||
"dpcId" : 0,
|
||||
"events" :[09:49:29.193][MCTX][I]server.cpp :348 : HandleRpcs exit. thread_no=[0].
|
||||
[09:49:30.193][MCTX][I]server.cpp :89 : Server Shutdown!
|
||||
[09:49:30.193][MCTX][I]mcRpcAsyncImpl.cpp :477 : mcRpcServerAsyncRun stop
|
||||
[09:49:30.275][MCTX][I]mctxTPImpl.cpp :433 : clear port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
|
||||
[info] profiling task 054b0685-1efd-4049-82d2-21af8938b66e
|
||||
[info] complete task 054b0685-1efd-4049-82d2-21af8938b66e
|
||||
stop server
|
||||
perf done, please check report file /opt/mcProfiler-ubuntu18.04/output20260717094852
|
||||
|
|
@ -0,0 +1,149 @@
|
|||
# opt_012 case 4 mcProfiler analysis (64 GiB C500)
|
||||
|
||||
## Scope
|
||||
|
||||
- Case: `equal_b1_s16384`
|
||||
- Kernel: `packed_kernel_dense_softmax_cleanup_v12_kernel`
|
||||
- Device: full MetaX C500, 64 GiB container
|
||||
- Collection: two per-kernel samples for each metric set, `--single-pass`
|
||||
- Reference: opt_007 case-4 mcTracer and mcProfiler reports
|
||||
|
||||
Both opt_012 samples are highly consistent. Total cycles differ by 0.32%, and
|
||||
the principal duty and memory metrics are effectively identical. The reports
|
||||
are suitable for bottleneck comparison.
|
||||
|
||||
## Timeline and resources
|
||||
|
||||
The opt_012 mcTracer result established:
|
||||
|
||||
```text
|
||||
device median: 46.012800 ms
|
||||
target-kernel share of the 20-launch device span: 99.994214%
|
||||
grid: (1024, 4, 1)
|
||||
block: (512, 1, 1)
|
||||
dynamic shared memory: 53,248 bytes/CTA
|
||||
registers: 254/thread
|
||||
private memory: 0
|
||||
MT-register occupancy: 49%
|
||||
shared-memory occupancy: 81%
|
||||
```
|
||||
|
||||
Case 4 is therefore device-kernel dominated. The shared-memory footprint still
|
||||
limits the kernel to one CTA per AP; there is no visible private-memory spill.
|
||||
|
||||
## Core counter comparison
|
||||
|
||||
The opt_012 values below are the mean of its two per-kernel samples. The opt_007
|
||||
values come from its corresponding per-kernel report.
|
||||
|
||||
| Metric | opt_007 | opt_012 | Change |
|
||||
|---|---:|---:|---:|
|
||||
| Total cycles | 55,369.05 K | 52,042.51 K | -6.01% |
|
||||
| AP busy duty | 95.64% | 99.29% | +3.65 pp |
|
||||
| AP MMA duty, core pass | 15.03% | 20.14% | +5.11 pp |
|
||||
| VL1 hit rate | 87.27% | 87.31% | +0.04 pp |
|
||||
| L2 hit rate | 95.24% | 97.50% | +2.26 pp |
|
||||
| Global read bytes | 711,921,664 | 310,777,856 | -56.35% |
|
||||
| Global write bytes | 134,219,776 | 134,217,728 | unchanged |
|
||||
| Shared-memory efficiency | 79.62% | 74.04% | -5.58 pp |
|
||||
| Achieved waves | 32,768 | 32,768 | unchanged |
|
||||
| Dispatched waves | 32,768 | 32,768 | unchanged |
|
||||
|
||||
The 6.01% cycle reduction agrees with mcTracer's 6.42% median-time reduction.
|
||||
This cross-tool agreement confirms that the speedup is real and is not caused by
|
||||
host timing or launch noise.
|
||||
|
||||
Global counter values collected with `--single-pass` are inferred from a limited
|
||||
event set. Their direction is useful, but the 56.35% global-read reduction should
|
||||
not be treated as an exact byte-saving proof without a non-single-pass control
|
||||
run. The unchanged 128-MiB write value is consistent with one output write.
|
||||
|
||||
## Pipeline duty comparison
|
||||
|
||||
| Pipeline metric | opt_007 | opt_012 | Change |
|
||||
|---|---:|---:|---:|
|
||||
| AP MTE duty / AP active | 38.86% | 38.31% | -0.55 pp |
|
||||
| AP STE duty / AP active | 1.70% | 1.09% | -0.61 pp |
|
||||
| AP MMA duty / AP active | 18.18% | 20.10% | +1.92 pp |
|
||||
| VLS duty / AP active | 0.0% | 0.0% | unchanged |
|
||||
| L2C duty / L2C active | 7.88% | 8.41% | +0.53 pp |
|
||||
|
||||
The second metric set confirms that opt_012 spends a larger share of active time
|
||||
issuing useful MMA work and less in STE work. MTE duty is essentially unchanged,
|
||||
so data movement remains a large part of the active pipeline mix. The reported
|
||||
zero VLS duty is not consistent with the nonzero VLS stall counters and should be
|
||||
treated as unsupported or uninformative on this tool version, not as proof that
|
||||
the kernel performs no vector loads/stores.
|
||||
|
||||
## Stall transition
|
||||
|
||||
Mean opt_012 stall values and opt_007 reference:
|
||||
|
||||
| Stall source | opt_007 cycles | opt_012 cycles | Relative change |
|
||||
|---|---:|---:|---:|
|
||||
| `wsm_stall` | 1,357,933,682 | 1,246,525,947 | -8.20% |
|
||||
| `vls_pipeline_stall` | 22,391,979 | 149,924,523 | +569.55% |
|
||||
| `vls_wdata_stall` | 2,097,152 | 2,097,152 | unchanged |
|
||||
|
||||
Distribution across these three reported sources:
|
||||
|
||||
| Stall share | opt_007 | opt_012 |
|
||||
|---|---:|---:|
|
||||
| `wsm_stall` | 98.23% | 89.13% |
|
||||
| `vls_pipeline_stall` | 1.62% | 10.72% |
|
||||
| `vls_wdata_stall` | 0.15% | 0.15% |
|
||||
|
||||
The original WSM constraint has improved in absolute terms but remains dominant.
|
||||
At the same time, VLS pipeline stalls increased substantially and shared-memory
|
||||
efficiency fell from 79.62% to 74.04%. The optimization shifted part of the
|
||||
bottleneck toward the vector/shared-memory access path.
|
||||
|
||||
## Bottleneck conclusion
|
||||
|
||||
1. Host launch, Python, synchronization, and inter-kernel gaps are eliminated as
|
||||
meaningful case-4 bottlenecks.
|
||||
2. opt_012's dense specialization and instruction cleanup reduce total work and
|
||||
increase MMA duty without changing launch geometry or residency.
|
||||
3. The kernel still has only one resident CTA/AP because it allocates 53,248
|
||||
bytes of dynamic shared memory. This limits latency hiding across the strict
|
||||
online-softmax dependency chain.
|
||||
4. Shared-memory efficiency regressed and VLS pipeline stalls grew. Shared-memory
|
||||
layout/access and vector-pipeline pressure are now the clearest immediate
|
||||
optimization targets.
|
||||
5. MMA duty improved to about 20%, but remains low relative to 99% AP busy duty.
|
||||
Most active time is still spent in movement, softmax, synchronization, or
|
||||
dependency waits rather than matrix instructions.
|
||||
|
||||
## Next controlled experiments
|
||||
|
||||
Run each experiment as a separate version and preserve opt_012 as the reference.
|
||||
|
||||
1. **Same 128x64 tile, shared-layout experiment.** Add only a MetaX-supported
|
||||
padding or swizzle to `q_shared` and `kv_shared`. Keep grid, threads, loop
|
||||
structure, and scheduling unchanged. Accept only if shared efficiency rises,
|
||||
VLS/WSM stalls fall, and correctness remains intact.
|
||||
2. **Inspect generated device code before changing stages.** Count shared-memory
|
||||
accesses and barriers in the KV loop, and identify the instructions associated
|
||||
with Q/K/V copies, score-to-probability conversion, and output rescaling.
|
||||
3. **Residency experiment, M64/N32/256 threads.** This is the realistic shape for
|
||||
reducing shared use below the two-CTA threshold. It doubles KV-loop iterations
|
||||
relative to N64, so it must be measured rather than assumed beneficial.
|
||||
4. **Reduce vector softmax traffic.** Look for a supported way to reduce fragment
|
||||
copies or repeated row-scale passes without increasing register spill. The
|
||||
acceptance signals are lower VLS pipeline stalls and unchanged private memory.
|
||||
5. **Do not enable two pipeline stages blindly.** Online softmax has loop-carried
|
||||
state, and the current shared footprint is already the residency limit. A
|
||||
second stage is justified only after a concrete double-buffered lifetime and
|
||||
resource calculation.
|
||||
|
||||
For every candidate, compare case-4 device time, resource metadata, shared
|
||||
efficiency, VLS/WSM stalls, and MMA duty against this report.
|
||||
|
||||
## Artifacts
|
||||
|
||||
```text
|
||||
results/tilelang_64g/opt_012_case4_mctracer_analysis_64g.md
|
||||
results/tilelang_64g/mcprofiler_opt012_case4_targeted/
|
||||
results/tilelang_64g/mcprofiler_opt012_case4_duty/
|
||||
results/tilelang_64g/opt_012_case4_mcprofiler_analysis_64g.md
|
||||
```
|
||||
|
|
@ -0,0 +1,385 @@
|
|||
* Serving Flask app 'Profiler' (lazy loading)
|
||||
* Environment: production
|
||||
WARNING: This is a development server. Do not use it in a production deployment.
|
||||
Use a production WSGI server instead.
|
||||
* Debug mode: off
|
||||
profiler http listen to 127.0.0.1:50123
|
||||
try to connect to server...
|
||||
connect to server success
|
||||
cmdline is: python tests/trace_tilelang_64g_case4.py
|
||||
casename is: opt012_case4_duty
|
||||
metrics is: ['Total Cycles', 'AP busy Duty', 'ISU stall cycles layout', 'AP MTE Duty ratio', 'AP STE Duty ratio', 'AP MMA Duty ratio', 'VLS Duty ratio', 'L2C Duty ratio']
|
||||
start a local task...
|
||||
[info] start a new perf exec thread
|
||||
[info] start perf_exec
|
||||
[info] new exec 577aafc1-d8d5-4d8f-9372-43de52e48907:('python tests/trace_tilelang_64g_case4.py', 'opt012_case4_duty', ['Total Cycles', 'AP busy Duty', 'ISU stall cycles layout', 'AP MTE Duty ratio', 'AP STE Duty ratio', 'AP MMA Duty ratio', 'VLS Duty ratio', 'L2C Duty ratio'])
|
||||
########## exec_id: 577aafc1-d8d5-4d8f-9372-43de52e48907 ##########
|
||||
[info] start generate batch files
|
||||
[info] output path is: /opt/mcProfiler-ubuntu18.04/output20260717095614
|
||||
[info] generate event_batch_0.json successfully!
|
||||
[info] generate batch files done
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
Loading tilelang libs from dev root: /data/tilelang-metax/build
|
||||
[09:56:16.212][MXS][I]| MACA_LAUNCH_BLOCKING | 1 | 0 |
|
||||
[09:56:16.212][MXS][W]| MCTX_TARGET_INIT | 1 | 0 |
|
||||
[09:56:16.215][MCR][I]mc_device.cpp :562 : Signal SIGUSR1 now has a custom handler to enable mctx at runtime!
|
||||
[09:56:16.215][MCR][I]mc_device.cpp :309 : mcSnapShotCtx: Attempting to open shared memory: /shm_snapshot_102181
|
||||
mctxAutoStart called!
|
||||
[09:56:16.215][MCTX][I]mctxTPImpl.cpp :201 : MACA configuration file path redirected from ${MACA_PATH}/etc to MACA_ETC_PATH
|
||||
[09:56:16.215][MCTX][I]mctxTPImpl.cpp :201 : MACA configuration file path redirected from ${MACA_PATH}/etc to MACA_ETC_PATH
|
||||
[09:56:16.215][MCTX][I]mctxTPImpl.cpp :300 : profiler target ENV file[/opt/mcProfiler-ubuntu18.04/output20260717095614/mcProfiler.json] exist!
|
||||
[09:56:16.216][MCR][I]mc_device.cpp :849 : Init enter!
|
||||
[09:56:16.216][MXC][I]MxcInit: entering
|
||||
[09:56:16.230][MXKW][I]topology.c :1358: [topology_sysfs_get_system_props]PlatformCpuNodes:2 ,num_sysfs_nodes:10
|
||||
[09:56:16.230][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/2/gpu_id,res: 0,errno:1
|
||||
[09:56:16.230][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 1,ret is 0,is_node_supported is 0
|
||||
[09:56:16.230][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/4/gpu_id,res: 0,errno:1
|
||||
[09:56:16.230][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 2,ret is 0,is_node_supported is 0
|
||||
[09:56:16.230][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/5/gpu_id,res: 0,errno:1
|
||||
[09:56:16.230][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 3,ret is 0,is_node_supported is 0
|
||||
[09:56:16.230][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/6/gpu_id,res: 0,errno:1
|
||||
[09:56:16.230][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 4,ret is 0,is_node_supported is 0
|
||||
[09:56:16.230][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/7/gpu_id,res: 0,errno:1
|
||||
[09:56:16.230][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 5,ret is 0,is_node_supported is 0
|
||||
[09:56:16.230][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/8/gpu_id,res: 0,errno:1
|
||||
[09:56:16.230][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 6,ret is 0,is_node_supported is 0
|
||||
[09:56:16.230][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/9/gpu_id,res: 0,errno:1
|
||||
[09:56:16.230][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 7,ret is 0,is_node_supported is 0
|
||||
[09:56:16.230][MXKW][I]topology.c :1460: [topology_sysfs_get_system_props]cpu total 2 nodes, map list:
|
||||
[09:56:16.230][MXKW][I]topology.c :1463: [topology_sysfs_get_system_props]map_cpu[0]:0
|
||||
[09:56:16.230][MXKW][I]topology.c :1463: [topology_sysfs_get_system_props]map_cpu[1]:1
|
||||
[09:56:16.230][MXKW][I]topology.c :1465: [topology_sysfs_get_system_props]gpu total 1 nodes, map list:
|
||||
[09:56:16.230][MXKW][I]topology.c :1468: [topology_sysfs_get_system_props]map_gpu[0]:3
|
||||
[09:56:16.230][MXKW][I]topology.c :1516: [topology_sysfs_init_gpu_page_size]XCORE_PAGE_SIZE: 0x200000, GPU_PAGE_SIZE: 0x200000
|
||||
[09:56:16.230][MXKW][I]topology.c :954 : [find_pci_info]create pci_info_array. It has 10 element. directory_name is 0.
|
||||
[09:56:16.230][MXKW][I]topology.c :962 : [find_pci_info]node 0 's pci info is not found
|
||||
[09:56:16.230][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/0/location file!
|
||||
[09:56:16.230][MXKW][I]topology.c :962 : [find_pci_info]node 0 's pci info is not found
|
||||
[09:56:16.230][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/0/location file!
|
||||
[09:56:16.230][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 0, isa_minor is 0, MXName is
|
||||
[09:56:16.230][MXKW][I]topology.c :962 : [find_pci_info]node 1 's pci info is not found
|
||||
[09:56:16.230][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/1/location file!
|
||||
[09:56:16.230][MXKW][I]topology.c :962 : [find_pci_info]node 1 's pci info is not found
|
||||
[09:56:16.230][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/1/location file!
|
||||
[09:56:16.230][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 0, isa_minor is 0, MXName is
|
||||
[09:56:16.230][MXKW][I]topology.c :962 : [find_pci_info]node 3 's pci info is not found
|
||||
[09:56:16.230][MXKW][I]topology.c :984 : [add_pci_info]added node 3 pci info , bdf_value:3840, domain_value:0
|
||||
[09:56:16.230][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 10, isa_minor is 0, MXName is
|
||||
[09:56:16.231][MXKW][I]fmm.c :2445: [init_mmap_apertures]Initialized unreserved SVM apertures: 0xa00000000 - 0x7fffffffffff
|
||||
[09:56:16.231][MXKW][I]topology.c :1358: [topology_sysfs_get_system_props]PlatformCpuNodes:2 ,num_sysfs_nodes:10
|
||||
[09:56:16.231][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/2/gpu_id,res: 0,errno:1
|
||||
[09:56:16.231][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 1,ret is 0,is_node_supported is 0
|
||||
[09:56:16.231][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/4/gpu_id,res: 0,errno:1
|
||||
[09:56:16.231][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 2,ret is 0,is_node_supported is 0
|
||||
[09:56:16.231][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/5/gpu_id,res: 0,errno:1
|
||||
[09:56:16.231][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 3,ret is 0,is_node_supported is 0
|
||||
[09:56:16.231][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/6/gpu_id,res: 0,errno:1
|
||||
[09:56:16.231][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 4,ret is 0,is_node_supported is 0
|
||||
[09:56:16.231][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/7/gpu_id,res: 0,errno:1
|
||||
[09:56:16.231][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 5,ret is 0,is_node_supported is 0
|
||||
[09:56:16.231][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/8/gpu_id,res: 0,errno:1
|
||||
[09:56:16.231][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 6,ret is 0,is_node_supported is 0
|
||||
[09:56:16.231][MXKW][I]topology.c :283 : [get_directory_socket_id]fscanf path:/sys/devices/virtual/mxcd/mxcd/layout/nodes/9/gpu_id,res: 0,errno:1
|
||||
[09:56:16.231][MXKW][W]topology.c :1424: [topology_sysfs_get_system_props]err_node_num is 7,ret is 0,is_node_supported is 0
|
||||
[09:56:16.231][MXKW][I]topology.c :1460: [topology_sysfs_get_system_props]cpu total 2 nodes, map list:
|
||||
[09:56:16.231][MXKW][I]topology.c :1463: [topology_sysfs_get_system_props]map_cpu[0]:0
|
||||
[09:56:16.231][MXKW][I]topology.c :1463: [topology_sysfs_get_system_props]map_cpu[1]:1
|
||||
[09:56:16.231][MXKW][I]topology.c :1465: [topology_sysfs_get_system_props]gpu total 1 nodes, map list:
|
||||
[09:56:16.231][MXKW][I]topology.c :1468: [topology_sysfs_get_system_props]map_gpu[0]:3
|
||||
[09:56:16.231][MXKW][I]topology.c :962 : [find_pci_info]node 0 's pci info is not found
|
||||
[09:56:16.231][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/0/location file!
|
||||
[09:56:16.231][MXKW][I]topology.c :962 : [find_pci_info]node 0 's pci info is not found
|
||||
[09:56:16.231][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/0/location file!
|
||||
[09:56:16.231][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 0, isa_minor is 0, MXName is
|
||||
[09:56:16.231][MXKW][I]topology.c :2548: [topology_take_snapshot]node0 direct_links route:
|
||||
[09:56:16.231][MXKW][I]topology.c :2562: [topology_take_snapshot]link0 type:2, 0--->1
|
||||
[09:56:16.231][MXKW][I]topology.c :2562: [topology_take_snapshot]link1 type:3, 0--->2
|
||||
[09:56:16.231][MXKW][I]topology.c :962 : [find_pci_info]node 1 's pci info is not found
|
||||
[09:56:16.231][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/1/location file!
|
||||
[09:56:16.231][MXKW][I]topology.c :962 : [find_pci_info]node 1 's pci info is not found
|
||||
[09:56:16.231][MXKW][I]topology.c :911 : [read_pci_info_from_file]can't open /sys/devices/virtual/mxcd/mxcd/layout/nodes/1/location file!
|
||||
[09:56:16.231][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 0, isa_minor is 0, MXName is
|
||||
[09:56:16.231][MXKW][I]topology.c :2548: [topology_take_snapshot]node1 direct_links route:
|
||||
[09:56:16.231][MXKW][I]topology.c :2562: [topology_take_snapshot]link0 type:2, 1--->0
|
||||
[09:56:16.231][MXKW][I]topology.c :2573: [topology_take_snapshot]node1 indirect_links route:
|
||||
[09:56:16.231][MXKW][I]topology.c :2597: [topology_take_snapshot]link1 type:105 1--->2
|
||||
[09:56:16.231][MXKW][I]topology.c :1944: [topology_sysfs_get_node_props]Device isa_major is 10, isa_minor is 0, MXName is
|
||||
[09:56:16.233][MXKW][I]topology.c :2548: [topology_take_snapshot]node2 direct_links route:
|
||||
[09:56:16.233][MXKW][I]topology.c :2562: [topology_take_snapshot]link0 type:3, 2--->0
|
||||
[09:56:16.233][MXKW][I]topology.c :2573: [topology_take_snapshot]node2 indirect_links route:
|
||||
[09:56:16.233][MXKW][I]topology.c :2597: [topology_take_snapshot]link1 type:105 2--->1
|
||||
[09:56:16.234][MXC][I]MxcInit: leaving
|
||||
[09:56:16.234][MXC][I]InitDma: Entering
|
||||
[09:56:16.237][MXC][I]InitDma: Leaving
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :2294: Device xcore1000 create, major:10, minor:0
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :2795: private_mem:436207616, info_.globalMemSize_:68283269120, gpu_single_alloc_percent:0.990000
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :2865: get registersPerMtreg_:64,info_dpcNum_:8,info_wavePerDpc_:416, wavefrontWidth_:64
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :3042: rmNum_:8, isuNum_:13, sl1Num_:7, vlsNum_:13, vl1Num_:13, l2cNum_:32, l2aNum_:4, l2xNum_:4
|
||||
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :1083: getQueueCnt ce queue bitmap:ffff
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :1091: getQueueCnt hwQueueNum:16
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :1153: agent_handle=0x56303b05b5a0,hardware queues number:16, queue count configuration: 8-4-4(H-N-L), highForMccl:0
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :2445: Device xcore1000 create fininshed!
|
||||
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :2294: Device xcore1000 create, major:10, minor:0
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :2795: private_mem:436207616, info_.globalMemSize_:68283269120, gpu_single_alloc_percent:0.990000
|
||||
[09:56:16.238][MCR][I]mx_device.cpp :2865: get registersPerMtreg_:64,info_dpcNum_:8,info_wavePerDpc_:416, wavefrontWidth_:64
|
||||
[09:56:16.239][MCR][I]mx_device.cpp :3042: rmNum_:8, isuNum_:13, sl1Num_:7, vlsNum_:13, vl1Num_:13, l2cNum_:32, l2aNum_:4, l2xNum_:4
|
||||
|
||||
[09:56:16.239][MCR][I]mx_device.cpp :1083: getQueueCnt ce queue bitmap:ffff
|
||||
[09:56:16.239][MCR][I]mx_device.cpp :1091: getQueueCnt hwQueueNum:16
|
||||
[09:56:16.239][MCR][I]mx_device.cpp :1153: agent_handle=0x56303b05b5a0,hardware queues number:16, queue count configuration: 8-4-4(H-N-L), highForMccl:0
|
||||
[09:56:16.239][MCR][I]mx_device.cpp :2445: Device xcore1000 create fininshed!
|
||||
|
||||
[09:56:16.239][MCR][I]mc_vpu_impl.cpp :80 : mcVpuLoad mxvpu lib version: 1.0.20260512_081136
|
||||
|
||||
[09:56:16.239][MCR][I]mc_vpu_impl.cpp :86 : mcVpuLoad load mxvpu lib path: /opt/maca-3.7.1/lib/libmxvpu.so
|
||||
|
||||
[09:56:16.239][MCC][I]AddDevice device_id:0, arch:xcore1000
|
||||
[09:56:16.239][MCC][I]Get compiler from environment variable MACA_CLANG_PATH: /opt/maca/mxgpu_llvm/bin.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[09:56:38.215][MXS][I] Node0 Node1
|
||||
[09:56:38.215][MXS][I]Gpu0 X
|
||||
[09:56:38.215][MXS][I]CpuBind X X
|
||||
get_mempolicy: Operation not permitted
|
||||
[09:56:38.215][MXS][I]MemBind X X
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[09:56:39.090][MXS][I]mcruntime 100% 0%
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[09:56:39.968][MXS][I]mccompiler 100% 0%
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[09:56:40.844][MXS][I]mxc-runtime 99% 1%
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[09:56:41.719][MXS][I]mxkw 100%
|
||||
[09:56:41.719][PTI][I]mcpti_profiler.cpp :618 : getModuleInfo:1,8,13,7,13,13,32,4,4
|
||||
|
||||
[09:56:41.721][MCR][I]mx_device.cpp :5792: agent_handle=0x56303b05b5a0
|
||||
[09:56:41.721][MXKW][I]queues.c :388 : [get_quantum_map_value]clockGrain 4
|
||||
[09:56:41.721][MXKW][I]queues.c :393 : [get_quantum_map_value]highQuantum 16 normal 8 low 4
|
||||
[09:56:41.722][MCR][I]mx_device.cpp :1446: created hardware queue 0x7f9d97a1e000(HWq=0x7f99a8001cd0) with size 1024 with priority 0(L), cooperative: 0, is_exclusive: 0
|
||||
[09:56:41.722][MCR][I]mx_device.cpp :1458: number of allocated hardware queues 0-0-1(H-N-L), number of users 0-0-1
|
||||
[09:56:41.725][MCR][W]mx_perf_counter.cpp :568 : Perf is not exclusive, can not be disabled!
|
||||
[09:56:41.725][MCR][W]mx_perf_counter.cpp :568 : Perf is not exclusive, can not be disabled!
|
||||
[09:56:41.726][MCTX][I]mctxTPImpl.cpp :186 : Automatically acquire a port.
|
||||
[09:56:41.726][MCTX][I]mcRpcAsyncImpl.cpp :421 : mcRpcServerAsyncInit 0x56303b04f790 0
|
||||
[09:56:41.726][MCTX][I]mcRpcAsyncImpl.cpp :500 : mcRpcAsyncServerTask start!
|
||||
[09:56:41.726][MCTX][I]mcRpcAsyncImpl.cpp :472 : mcRpcServerAsyncRun start
|
||||
[09:56:41.727][MCTX][I]server.cpp :69 : Server thread num: 1, listening on: 39945.
|
||||
[info] MctxStreamProfilerCountDataGet Start
|
||||
[info] MctxStreamProfilerCountDataGet Loop
|
||||
mcToolsExtPid_UsedPort:102181-39945
|
||||
[09:56:45.442][MXC][I]raw dispatch ts:[252068959726940,252068959789928]
|
||||
[09:56:45.442][MXC][I]dispatch ts translate ret [1784282205441364895, 1784282205441994772]
|
||||
[09:56:45.442][MXC][I]raw dispatch ts:[252068959809647,252068959818842]
|
||||
[09:56:45.442][MXC][I]dispatch ts translate ret [1784282205442191926, 1784282205442283852]
|
||||
[09:56:45.442][MXC][I]raw dispatch ts:[252068959840003,252068959849325]
|
||||
[09:56:45.442][MXC][I]dispatch ts translate ret [1784282205442495500, 1784282205442588741]
|
||||
[09:56:45.442][MCR][I]mc_memory.cpp :1671: copy strategy - H2DCpuCopy:0,D2HCpuCopy:0,cpuCopy:0,d2dCrossDevices:0,pageableHToDAsync:0
|
||||
[09:56:45.442][MCR][I]mc_memory.cpp :1672: copy strategy - isAsync:0,waitComplete:1
|
||||
[09:56:45.442][MXC][I]GpuDevice 0x3b05b5a0 enable profiling, sdm_counter:202851652513868, sys_time:1784282205442923326
|
||||
[09:56:45.442][MXC][I]raw async_copy ts:[202851652533034,202851652535620]
|
||||
[09:56:45.442][MXC][I]sdma_ts_.sdma_counter:202851652513868
|
||||
[09:56:45.442][MXC][I]DNOCfreq:900000000, async_copy ts translate ret [1784282205442944621, 1784282205442947494]
|
||||
[09:56:49.540][MXC][I]raw dispatch ts:[252069365081844,252069369684876]
|
||||
[09:56:49.540][MXC][I]dispatch ts translate ret [1784282209494897644, 1784282209540927779]
|
||||
[info] MctxStreamProfilerCountDataGet Loop
|
||||
[09:56:49.587][MXC][I]raw dispatch ts:[252069369737563,252069374342478]
|
||||
[09:56:49.587][MXC][I]dispatch ts translate ret [1784282209541454647, 1784282209587503617]
|
||||
[09:56:49.633][MXC][I]raw dispatch ts:[252069374371453,252069378974857]
|
||||
[09:56:49.633][MXC][I]dispatch ts translate ret [1784282209587793366, 1784282209633827225]
|
||||
[09:56:49.679][MXC][I]raw dispatch ts:[252069378986174,252069383588798]
|
||||
[09:56:49.680][MXC][I]dispatch ts translate ret [1784282209633940394, 1784282209679966428]
|
||||
[09:56:49.726][MXC][I]raw dispatch ts:[252069383600189,252069388203670]
|
||||
[09:56:49.726][MXC][I]dispatch ts translate ret [1784282209680080337, 1784282209726114974]
|
||||
[09:56:49.772][MXC][I]raw dispatch ts:[252069388216463,252069392820065]
|
||||
[09:56:49.772][MXC][I]dispatch ts translate ret [1784282209726242903, 1784282209772278738]
|
||||
[info] MctxStreamProfilerCountDataGet Loop
|
||||
[09:56:49.818][MXC][I]raw dispatch ts:[252069392830979,252069397432392]
|
||||
[09:56:49.818][MXC][I]dispatch ts translate ret [1784282209772387877, 1784282209818401831]
|
||||
[09:56:49.864][MXC][I]raw dispatch ts:[252069397446703,252069402047915]
|
||||
[09:56:49.864][MXC][I]dispatch ts translate ret [1784282209818544940, 1784282209864556866]
|
||||
[09:56:49.910][MXC][I]raw dispatch ts:[252069402058757,252069406659014]
|
||||
[09:56:49.910][MXC][I]dispatch ts translate ret [1784282209864665285, 1784282209910667673]
|
||||
[09:56:49.956][MXC][I]raw dispatch ts:[252069406669311,252069411269645]
|
||||
[09:56:49.956][MXC][I]dispatch ts translate ret [1784282209910770642, 1784282209956773807]
|
||||
[09:56:50.002][MXC][I]raw dispatch ts:[252069411280116,252069415882571]
|
||||
[09:56:50.002][MXC][I]dispatch ts translate ret [1784282209956878516, 1784282210002902883]
|
||||
TRACE_REGION_BEGIN repeats=20
|
||||
[09:56:50.049][MXC][I]raw dispatch ts:[252069415900942,252069420504316]
|
||||
[09:56:50.049][MXC][I]dispatch ts translate ret [1784282210003086592, 1784282210049120137]
|
||||
[09:56:50.095][MXC][I]raw dispatch ts:[252069420515347,252069425116419]
|
||||
[09:56:50.095][MXC][I]dispatch ts translate ret [1784282210049230446, 1784282210095240997]
|
||||
[09:56:50.141][MXC][I]raw dispatch ts:[252069425130231,252069429732109]
|
||||
[09:56:50.141][MXC][I]dispatch ts translate ret [1784282210095379116, 1784282210141397698]
|
||||
[09:56:50.187][MXC][I]raw dispatch ts:[252069429741220,252069434342012]
|
||||
[09:56:50.187][MXC][I]dispatch ts translate ret [1784282210141488807, 1784282210187496538]
|
||||
[09:56:50.233][MXC][I]raw dispatch ts:[252069434351834,252069438952531]
|
||||
[09:56:50.233][MXC][I]dispatch ts translate ret [1784282210187594757, 1784282210233601547]
|
||||
[09:56:50.279][MXC][I]raw dispatch ts:[252069438963609,252069443567746]
|
||||
[09:56:50.279][MXC][I]dispatch ts translate ret [1784282210233712326, 1784282210279753521]
|
||||
[09:56:50.325][MXC][I]raw dispatch ts:[252069443577738,252069448180989]
|
||||
[09:56:50.325][MXC][I]dispatch ts translate ret [1784282210279853441, 1784282210325885985]
|
||||
[09:56:50.372][MXC][I]raw dispatch ts:[252069448191227,252069452794172]
|
||||
[09:56:50.372][MXC][I]dispatch ts translate ret [1784282210325988365, 1784282210372017409]
|
||||
[09:56:50.418][MXC][I]raw dispatch ts:[252069452804012,252069457405310]
|
||||
[09:56:50.418][MXC][I]dispatch ts translate ret [1784282210372115808, 1784282210418128606]
|
||||
[09:56:50.464][MXC][I]raw dispatch ts:[252069457415442,252069462018434]
|
||||
[09:56:50.464][MXC][I]dispatch ts translate ret [1784282210418229925, 1784282210464259660]
|
||||
[09:56:50.510][MXC][I]raw dispatch ts:[252069462030134,252069466630818]
|
||||
[09:56:50.510][MXC][I]dispatch ts translate ret [1784282210464376659, 1784282210510383314]
|
||||
[09:56:50.556][MXC][I]raw dispatch ts:[252069466642079,252069471243554]
|
||||
[09:56:50.556][MXC][I]dispatch ts translate ret [1784282210510495923, 1784282210556510489]
|
||||
[09:56:50.602][MXC][I]raw dispatch ts:[252069471256288,252069475855218]
|
||||
[09:56:50.602][MXC][I]dispatch ts translate ret [1784282210556637828, 1784282210602626943]
|
||||
[09:56:50.648][MXC][I]raw dispatch ts:[252069475865856,252069480467022]
|
||||
[09:56:50.648][MXC][I]dispatch ts translate ret [1784282210602733322, 1784282210648744798]
|
||||
[09:56:50.694][MXC][I]raw dispatch ts:[252069480477167,252069485078109]
|
||||
[09:56:50.694][MXC][I]dispatch ts translate ret [1784282210648846247, 1784282210694855495]
|
||||
[09:56:50.741][MXC][I]raw dispatch ts:[252069485090275,252069489694434]
|
||||
[09:56:50.741][MXC][I]dispatch ts translate ret [1784282210694977154, 1784282210741018537]
|
||||
[09:56:50.787][MXC][I]raw dispatch ts:[252069489706600,252069494307902]
|
||||
[09:56:50.787][MXC][I]dispatch ts translate ret [1784282210741140196, 1784282210787153051]/tmp/_MEI0UzLWv/phttp/http_server.py:89: UserWarning:
|
||||
|
||||
The 'environ['werkzeug.server.shutdown']' function is deprecated and will be removed in Werkzeug 2.1.
|
||||
|
||||
|
||||
[09:56:50.833][MXC][I]raw dispatch ts:[252069494318798,252069498920962]
|
||||
[09:56:50.833][MXC][I]dispatch ts translate ret [1784282210787262010, 1784282210833283467]
|
||||
[09:56:50.879][MXC][I]raw dispatch ts:[252069498931922,252069503531948]
|
||||
[09:56:50.879][MXC][I]dispatch ts translate ret [1784282210833393066, 1784282210879393141]
|
||||
[09:56:50.925][MXC][I]raw dispatch ts:[252069503541845,252069508140374]
|
||||
[09:56:50.925][MXC][I]dispatch ts translate ret [1784282210879492110, 1784282210925477217]
|
||||
TRACE_REGION_END
|
||||
[09:56:50.926][MXC][I]raw dispatch ts:[252069508243987,252069508262780]
|
||||
[09:56:50.926][MXC][I]dispatch ts translate ret [1784282210926513342, 1784282210926701272]
|
||||
[09:56:50.928][MXC][I]raw dispatch ts:[252069508400043,252069508414568]
|
||||
[09:56:50.928][MXC][I]dispatch ts translate ret [1784282210928073893, 1784282210928219142]
|
||||
[09:56:50.928][MXC][I]raw dispatch ts:[252069508435916,252069508450838]
|
||||
[09:56:50.928][MXC][I]dispatch ts translate ret [1784282210928432612, 1784282210928581824]
|
||||
[09:56:50.930][MXC][I]raw dispatch ts:[252069508617779,252069508632234]
|
||||
[09:56:50.930][MXC][I]dispatch ts translate ret [1784282210930251225, 1784282210930395775]
|
||||
[09:56:50.930][MXC][I]raw dispatch ts:[252069508669861,252069508670383]
|
||||
[09:56:50.930][MXC][I]dispatch ts translate ret [1784282210930772010, 1784282210930777230]
|
||||
[09:56:50.933][MXC][I]raw dispatch ts:[252069508892189,252069508900152]
|
||||
[09:56:50.933][MXC][I]dispatch ts translate ret [1784282210932995306, 1784282210933074936]
|
||||
[09:56:50.933][MCR][I]mc_memory.cpp :1671: copy strategy - H2DCpuCopy:0,D2HCpuCopy:0,cpuCopy:0,d2dCrossDevices:0,pageableHToDAsync:0
|
||||
[09:56:50.933][MCR][I]mc_memory.cpp :1672: copy strategy - isAsync:0,waitComplete:0
|
||||
[09:56:50.933][MXC][I]raw async_copy ts:[202858240993819,202858240995704]
|
||||
[09:56:50.933][MXC][I]sdma_ts_.sdma_counter:202851652513868
|
||||
[09:56:50.933][MXC][I]DNOCfreq:900000000, async_copy ts translate ret [1784282212763456604, 1784282212763458698]
|
||||
[09:56:50.933][MCR][W]mx_command.cpp :1883: HW exec timestamp will be adjusted!!!
|
||||
[09:56:50.933][MCR][W]mx_command.cpp :1885: floor:1784282210933289811, start:1784282212763456512, end:1784282212763458816, ceiling:1784282210933311481
|
||||
TILELANG_64G_CASE4_TRACE_PASS
|
||||
[09:56:51.159][MCR][I]mc_device.cpp :1002: tearDown enter!
|
||||
|
||||
[09:56:51.159][MCR][I]mx_device.cpp :1578: deleting hardware queue 0x7f9d97a1e000 with refCount 0
|
||||
[09:56:51.159][MCR][I]mx_device.cpp :1601: number of allocated hardware queues 0-0-0(H-N-L), number of users 0-0-0
|
||||
[09:56:51.159][MXC][I]~GpuDevice: Entering
|
||||
[09:56:51.160][MCTX][I]mcRpcAsyncImpl.cpp :550 : callback mcRpcAsyncDone!
|
||||
[09:56:51.160][PTI][I]mcpti_tracer.hpp :309 : stopProcessThread
|
||||
[09:56:51.160][MCTX][I]mcRpcAsyncImpl.cpp :491 : mcRpcServerAsync TearDown!
|
||||
[09:56:51.160][MCTX][I]server.cpp :116 : Server will tearDown[3].
|
||||
[09:56:51.171][MCTX][I]mcRpcAsyncImpl.cpp :250 : ProfilerCountDataGet the last msg:[
|
||||
{
|
||||
"CE" : [
|
||||
{
|
||||
"blockId" : 0,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"2" : 175719836902556,
|
||||
"4" : 450072933279
|
||||
}
|
||||
}
|
||||
],
|
||||
"ISU" : [
|
||||
{
|
||||
"blockId" : 0,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"1" : 1212512643443,
|
||||
"10" : 1786443662588,
|
||||
"40" : 1347817262993,
|
||||
"41" : 1246633206573,
|
||||
"79" : 7261267993368,
|
||||
"81" : 9324992117274,
|
||||
"82" : 1362363000338
|
||||
}
|
||||
},
|
||||
{
|
||||
"blockId" : 1,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"1" : 1212538983697,
|
||||
"101" : 1647186773391,
|
||||
"43" : 1256805977455,
|
||||
"52" : 19721579495299,
|
||||
"80" : 16744813480297,
|
||||
"89" : 2423902457246
|
||||
}
|
||||
}
|
||||
],
|
||||
"L2C" : [
|
||||
[info] MctxStreamProfilerCountDataGet Loop
|
||||
{
|
||||
"blockId" : 0,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"1" : 744790559108,
|
||||
"74" : 24301834034
|
||||
}
|
||||
}
|
||||
],
|
||||
"VLS" : [
|
||||
{
|
||||
"blockId" : 0,
|
||||
"dpcId" : 0,
|
||||
"events" : {
|
||||
"23" : 0
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
[09:56:51.182][MCTX][I]server.cpp :348 : HandleRpcs exit. thread_no=[0].
|
||||
[09:56:52.182][MCTX][I]server.cpp :89 : Server Shutdown!
|
||||
[09:56:52.182][MCTX][I]mcRpcAsyncImpl.cpp :477 : mcRpcServerAsyncRun stop
|
||||
[09:56:52.261][MCTX][I]mctxTPImpl.cpp :433 : clear port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
[warning] get_srvrpc_port get rpc server ports failed: cannot retrieve rpc server port.
|
||||
|
||||
[info] profiling task 577aafc1-d8d5-4d8f-9372-43de52e48907
|
||||
[info] complete task 577aafc1-d8d5-4d8f-9372-43de52e48907
|
||||
stop server
|
||||
perf done, please check report file /opt/mcProfiler-ubuntu18.04/output20260717095614
|
||||
|
|
@ -0,0 +1,147 @@
|
|||
# opt_012 case 4 mcTracer analysis on the 64 GiB C500
|
||||
|
||||
## Trace setup
|
||||
|
||||
```text
|
||||
case: equal_b1_s16384
|
||||
kernel: packed_kernel_dense_softmax_cleanup_v12_kernel
|
||||
first launch: 1
|
||||
warmup launches: 10
|
||||
analyzed consecutive launches: 20
|
||||
```
|
||||
|
||||
The process loads PyTorch and the TileLang kernel only. The trace harness loads
|
||||
`tilelang/run_kernel.py`, whose SHA-256 is identical to
|
||||
`tilelang/opt_012_dense_softmax_cleanup.py`. The import alias in the harness still
|
||||
contains `opt007`, but that alias does not select the kernel and does not affect
|
||||
the trace.
|
||||
|
||||
## Collection result
|
||||
|
||||
```text
|
||||
trace events: 419
|
||||
target kernel events: 31 = 1 first launch + 10 warmup + 20 analyzed
|
||||
trace file size: approximately 115 KiB
|
||||
```
|
||||
|
||||
The final 20 target launches contain no device memcpy or unrelated GPU kernel.
|
||||
They are asynchronously enqueued and followed by one device synchronization.
|
||||
Consequently, the increasing `queue_ts -> start_ts` of later launches is time
|
||||
spent behind earlier kernels in the same queue, not per-launch scheduling
|
||||
overhead.
|
||||
|
||||
## Kernel launch metadata
|
||||
|
||||
```text
|
||||
grid: (1024, 4, 1)
|
||||
block: (512, 1, 1)
|
||||
dynamic shared memory: 53248 bytes
|
||||
static shared memory: 0 bytes
|
||||
private memory per thread: 0 bytes
|
||||
private memory total: 0 bytes
|
||||
registers per thread: 254
|
||||
MT register occupancy: 49%
|
||||
shared-memory occupancy: 81%
|
||||
hardware queue: 2
|
||||
recompiled: false
|
||||
```
|
||||
|
||||
There is no private-memory spill allocation visible in mcTracer. Dynamic shared
|
||||
memory remains large enough to prevent two 53,248-byte CTAs from fitting in a
|
||||
64-KiB WSM allocation.
|
||||
|
||||
## Timeline statistics
|
||||
|
||||
For the final 20 consecutive target kernels:
|
||||
|
||||
```text
|
||||
duration minimum: 45.988096 ms
|
||||
duration P10: 45.994522 ms
|
||||
duration median: 46.012800 ms
|
||||
duration mean: 46.013210 ms
|
||||
duration P90: 46.036941 ms
|
||||
duration maximum: 46.040576 ms
|
||||
population stdev: 0.015535 ms
|
||||
coefficient of variation: 0.0338%
|
||||
```
|
||||
|
||||
The run is extremely stable. Across the complete 20-launch device interval:
|
||||
|
||||
```text
|
||||
sum of kernel durations: 920.264192 ms
|
||||
first-start to last-end span: 920.317440 ms
|
||||
sum of 19 inter-kernel gaps: 0.053248 ms
|
||||
target-kernel share of the span: 99.994214%
|
||||
```
|
||||
|
||||
The matched runtime `mcLaunchKernel` calls for the same 20 kernels take:
|
||||
|
||||
```text
|
||||
minimum: 4.120 us
|
||||
median: 4.592 us
|
||||
mean: 5.275 us
|
||||
maximum: 17.877 us
|
||||
sum: 105.507 us
|
||||
```
|
||||
|
||||
Queue submission itself has a median latency of 1.340 us. Host launch overhead,
|
||||
Python overhead, synchronization, and device bubbles are therefore negligible
|
||||
relative to one approximately 46-ms kernel.
|
||||
|
||||
## Comparison with opt_007
|
||||
|
||||
Using the same final-20-launch mcTracer method:
|
||||
|
||||
| Metric | opt_007 | opt_012 | Change |
|
||||
|---|---:|---:|---:|
|
||||
| Kernel median | 49.170432 ms | 46.012800 ms | -6.42% |
|
||||
| Kernel speed | 1.000x | 1.0686x | +6.86% |
|
||||
| Dynamic shared memory | 53,248 B | 53,248 B | unchanged |
|
||||
| Registers/thread | 255 | 254 | -1 |
|
||||
| MT register occupancy | 49% | 49% | unchanged |
|
||||
| Shared-memory occupancy | 81% | 81% | unchanged |
|
||||
| Private memory | 0 B | 0 B | unchanged |
|
||||
|
||||
The improvement did not come from higher reported occupancy, smaller shared
|
||||
memory, fewer CTAs, or removal of a spill. It is consistent with opt_012's
|
||||
within-kernel work reduction and scheduling changes: dense equal-length
|
||||
specialization, reverse dense tile order, fully-visible causal tile cleanup, and
|
||||
softmax cleanup. mcTracer measures their aggregate result but cannot attribute
|
||||
time to individual instructions.
|
||||
|
||||
## Bottleneck conclusion
|
||||
|
||||
1. Case 4 is entirely device-kernel dominated. Optimizing wrappers, launch count,
|
||||
or host synchronization cannot materially improve this testcase.
|
||||
2. The current kernel still runs one CTA per AP due to its 53,248-byte dynamic
|
||||
shared-memory footprint. Registers changed only from 255 to 254, leaving the
|
||||
reported 49% MT-register occupancy unchanged.
|
||||
3. opt_012 improves the device median by 3.157632 ms over opt_007 while preserving
|
||||
the same launch geometry and effective resource limits. This validates the
|
||||
instruction/work-reduction direction, but does not remove the original
|
||||
shared-memory residency constraint.
|
||||
4. mcTracer cannot distinguish MMA under-utilization, shared-memory conflicts,
|
||||
barriers, instruction dependencies, or cache/bandwidth limits. A filtered
|
||||
opt_012 mcProfiler run is required to compare those counters directly with the
|
||||
existing opt_007 profile.
|
||||
|
||||
## Next profiling comparison
|
||||
|
||||
Collect the same opt_012 target-kernel mcProfiler metrics used for opt_007:
|
||||
|
||||
- AP busy, AP MMA/MTE/STE duty;
|
||||
- shared-memory access efficiency and ISU stall breakdown;
|
||||
- global read/write bytes, VL1 and L2 hit rates;
|
||||
- achieved/dispatched waves and runtime resource metadata.
|
||||
|
||||
The decisive comparison is whether opt_012's 6.42% time reduction came from
|
||||
higher MMA duty, lower `wsm_stall`, fewer memory transactions, or fewer non-MMA
|
||||
instructions. Use the per-kernel report rather than the aggregate process report.
|
||||
|
||||
## Artifacts
|
||||
|
||||
```text
|
||||
tests/trace_tilelang_64g_case4.py
|
||||
results/tilelang_64g/mctracer_opt012_case4/tracer_out-98428.json
|
||||
results/tilelang_64g/opt_012_case4_mctracer_analysis_64g.md
|
||||
```
|
||||
|
|
@ -0,0 +1,168 @@
|
|||
#include <tl_templates/maca/gemm.h>
|
||||
#include <tl_templates/maca/copy.h>
|
||||
#include <tl_templates/maca/reduce.h>
|
||||
#include <tl_templates/maca/intrin.h>
|
||||
#include <tl_templates/maca/atomic.h>
|
||||
#include <tl_templates/maca/threadblock_swizzle.h>
|
||||
#include <tl_templates/maca/debug.h>
|
||||
|
||||
extern "C" __global__ void packed_kernel_v_prefetch_after_qk_v15_kernel(const bfloat16_t* __restrict__ k, const int* __restrict__ kv_indptr, bfloat16_t* __restrict__ output, const bfloat16_t* __restrict__ q, const bfloat16_t* __restrict__ v);
|
||||
extern "C" __global__ void __launch_bounds__(512, 1) packed_kernel_v_prefetch_after_qk_v15_kernel(const bfloat16_t* __restrict__ k, const int* __restrict__ kv_indptr, bfloat16_t* __restrict__ output, const bfloat16_t* __restrict__ q, const bfloat16_t* __restrict__ v) {
|
||||
extern __shared__ __align__(1024) uchar buf_dyn_shmem[];
|
||||
float output_accum[32];
|
||||
float row_denom[1];
|
||||
float row_max[1];
|
||||
float row_scale[1];
|
||||
float scores[16];
|
||||
float row_max_prev[1];
|
||||
bfloat16_t probs[16];
|
||||
bfloat16_t output_local_cast[4];
|
||||
int copy_kv_len = kv_indptr[1];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + ((((((((((int)threadIdx.x) & 15) >> 3) * 8192) + (i * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8))) = *(uint4*)(q + ((((((i * 16384) + ((((int)threadIdx.x) >> 7) * 4096)) + (((int)blockIdx.y) * 1024)) + ((((int)threadIdx.x) & 127) * 8)) + 67043328) - (((int)blockIdx.x) * 65536)));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i_1 = 0; i_1 < 8; ++i_1) {
|
||||
float broadcast_var = 0x0p+0f/*0.000000e+00*/;
|
||||
*(float4*)(output_accum + (i_1 * 4)) = make_float4(broadcast_var, broadcast_var, broadcast_var, broadcast_var);
|
||||
}
|
||||
row_denom[0] = 0x0p+0f/*0.000000e+00*/;
|
||||
row_max[0] = -MACART_INF_F;
|
||||
for (int kv_tile = 0; kv_tile < ((16447 - (((int)blockIdx.x) * 16)) >> 6); ++kv_tile) {
|
||||
__syncthreads();
|
||||
if (((kv_tile * 64) + 64) <= copy_kv_len) {
|
||||
#pragma unroll
|
||||
for (int i_2 = 0; i_2 < 2; ++i_2) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_2 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(k + (((((kv_tile * 32768) + (i_2 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_3 = 0; i_3 < 2; ++i_3) {
|
||||
if ((((kv_tile * 64) + (i_3 * 32)) + (((int)threadIdx.x) >> 4)) < copy_kv_len) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_3 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(k + (((((kv_tile * 32768) + (i_3 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
} else {
|
||||
bfloat16_t broadcast_var_1 = bfloat16_t(0x0p+0f/*0.000000e+00*/);
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_3 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = make_uint4(__pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1));
|
||||
}
|
||||
}
|
||||
}
|
||||
if (kv_tile < ((16369 - (((int)blockIdx.x) * 16)) >> 6)) {
|
||||
#pragma unroll
|
||||
for (int i_4 = 0; i_4 < 4; ++i_4) {
|
||||
float broadcast_var_2 = 0x0p+0f/*0.000000e+00*/;
|
||||
*(float4*)(scores + (i_4 * 4)) = make_float4(broadcast_var_2, broadcast_var_2, broadcast_var_2, broadcast_var_2);
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_5 = 0; i_5 < 16; ++i_5) {
|
||||
float condval;
|
||||
if ((((((kv_tile * 64) + ((i_5 >> 2) * 16)) + (((((int)threadIdx.x) & 63) >> 4) * 4)) + (i_5 & 3)) <= (((((((int)threadIdx.x) >> 6) * 2) + ((((int)threadIdx.x) & 15) >> 3)) + 16368) - (((int)blockIdx.x) * 16)))) {
|
||||
condval = 0x0p+0f/*0.000000e+00*/;
|
||||
} else {
|
||||
condval = -MACART_INF_F;
|
||||
}
|
||||
scores[i_5] = condval;
|
||||
}
|
||||
}
|
||||
bfloat16_t A_local[4];
|
||||
bfloat16_t B_local[16];
|
||||
__syncthreads();
|
||||
for (int ki = 0; ki < 8; ++ki) {
|
||||
*(uint2*)(A_local + 0) = *(uint2*)(((bfloat16_t*)buf_dyn_shmem) + ((((((((ki >> 2) * 8192) + ((((int)threadIdx.x) >> 6) * 1024)) + ((((int)threadIdx.x) & 15) * 64)) + (((((((int)threadIdx.x) & 7) >> 2) + ((ki & 3) >> 1)) & 1) * 32)) + (((((((int)threadIdx.x) & 3) >> 1) + (ki & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 63) >> 5) + (((int)threadIdx.x) & 1)) & 1) * 8)) + (((((int)threadIdx.x) & 31) >> 4) * 4)));
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
*(uint2*)(B_local + (j * 4)) = *(uint2*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((ki >> 2) * 4096) + (j * 1024)) + ((((int)threadIdx.x) & 15) * 64)) + (((((((int)threadIdx.x) & 7) >> 2) + ((ki & 3) >> 1)) & 1) * 32)) + (((((((int)threadIdx.x) & 3) >> 1) + (ki & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 63) >> 5) + (((int)threadIdx.x) & 1)) & 1) * 8)) + (((((int)threadIdx.x) & 31) >> 4) * 4)) + 16384));
|
||||
}
|
||||
for (int j_1 = 0; j_1 < 4; ++j_1) {
|
||||
{
|
||||
*(((float32x4*)scores) + j_1) = __builtin_mxc_mma_16x16x16bf16(*(((bfloat16x4_vec*)B_local) + j_1),
|
||||
*(((bfloat16x4_vec*)A_local) + 0),
|
||||
*(((float32x4*)scores) + j_1));
|
||||
};
|
||||
}
|
||||
}
|
||||
row_max_prev[0] = -MACART_INF_F;
|
||||
#pragma unroll
|
||||
for (int rv = 0; rv < 16; ++rv) {
|
||||
row_max_prev[0] = max(row_max_prev[0], scores[(((rv & 3) * 4) + (rv >> 2))]);
|
||||
}
|
||||
__syncthreads();
|
||||
row_max_prev[0] = tl::AllReduce<tl::MaxOp, 64, 16, 0>::run(row_max_prev[0], (&(((float*)buf_dyn_shmem)[12800])));
|
||||
__syncthreads();
|
||||
if (((kv_tile * 64) + 64) <= copy_kv_len) {
|
||||
#pragma unroll
|
||||
for (int i_6 = 0; i_6 < 2; ++i_6) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_6 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(v + (((((kv_tile * 32768) + (i_6 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_7 = 0; i_7 < 2; ++i_7) {
|
||||
if ((((kv_tile * 64) + (i_7 * 32)) + (((int)threadIdx.x) >> 4)) < copy_kv_len) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_7 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(v + (((((kv_tile * 32768) + (i_7 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
} else {
|
||||
bfloat16_t broadcast_var_3 = bfloat16_t(0x0p+0f/*0.000000e+00*/);
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_7 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = make_uint4(__pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3));
|
||||
}
|
||||
}
|
||||
}
|
||||
row_scale[0] = exp2f(((row_max[0] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/) - (max(row_max[0], row_max_prev[0]) * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/)));
|
||||
row_max[0] = max(row_max[0], row_max_prev[0]);
|
||||
#pragma unroll
|
||||
for (int i_8 = 0; i_8 < 16; ++i_8) {
|
||||
scores[i_8] = exp2f(((scores[i_8] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/) - (row_max[0] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/)));
|
||||
}
|
||||
row_max_prev[0] = 0x0p+0f/*0.000000e+00*/;
|
||||
#pragma unroll
|
||||
for (int rv_1 = 0; rv_1 < 16; ++rv_1) {
|
||||
row_max_prev[0] = (row_max_prev[0] + scores[(((rv_1 & 3) * 4) + (rv_1 >> 2))]);
|
||||
}
|
||||
__syncthreads();
|
||||
row_max_prev[0] = tl::AllReduce<tl::SumOp, 64, 16, 0>::run(row_max_prev[0], (&(((float*)buf_dyn_shmem)[12288])));
|
||||
row_denom[0] = ((row_denom[0] * row_scale[0]) + row_max_prev[0]);
|
||||
#pragma unroll
|
||||
for (int i_9 = 0; i_9 < 4; ++i_9) {
|
||||
uint2 __1;
|
||||
float4 v_ = *(float4*)(scores + (i_9 * 4));
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__1))[0] = __float22bfloat162_rn(((float2*)(&v_))[0]);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__1))[1] = __float22bfloat162_rn(((float2*)(&v_))[1]);
|
||||
*(uint2*)(probs + (i_9 * 4)) = __1;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i_10 = 0; i_10 < 32; ++i_10) {
|
||||
output_accum[i_10] = (output_accum[i_10] * row_scale[0]);
|
||||
}
|
||||
bfloat16_t B_local_1[32];
|
||||
__syncthreads();
|
||||
for (int ki_1 = 0; ki_1 < 4; ++ki_1) {
|
||||
for (int j_2 = 0; j_2 < 8; ++j_2) {
|
||||
for (int local_id = 0; local_id < 4; ++local_id) {
|
||||
B_local_1[((j_2 * 4) + local_id)] = ((bfloat16_t*)buf_dyn_shmem)[((((((((((j_2 >> 2) * 4096) + (ki_1 * 1024)) + (((((int)threadIdx.x) & 63) >> 4) * 256)) + (local_id * 64)) + (((((((int)threadIdx.x) & 31) >> 4) + ((j_2 & 3) >> 1)) & 1) * 32)) + ((((local_id >> 1) + (j_2 & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 15) >> 3) + (local_id & 1)) & 1) * 8)) + (((int)threadIdx.x) & 7)) + 16384)];
|
||||
}
|
||||
}
|
||||
for (int j_3 = 0; j_3 < 8; ++j_3) {
|
||||
{
|
||||
*(((float32x4*)output_accum) + j_3) = __builtin_mxc_mma_16x16x16bf16(*(((bfloat16x4_vec*)B_local_1) + j_3),
|
||||
*(((bfloat16x4_vec*)probs) + ki_1),
|
||||
*(((float32x4*)output_accum) + j_3));
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
row_scale[0] = (0x1p+0f/*1.000000e+00*/ / row_denom[0]);
|
||||
#pragma unroll
|
||||
for (int i_11 = 0; i_11 < 8; ++i_11) {
|
||||
uint2 __2;
|
||||
float4 __3;
|
||||
float4 v__1 = *(float4*)(output_accum + (i_11 * 4));
|
||||
float4 v__2 = make_float4(row_scale[0], row_scale[0], row_scale[0], row_scale[0]);
|
||||
__3.x = (v__1.x*v__2.x);
|
||||
__3.y = (v__1.y*v__2.y);
|
||||
__3.z = (v__1.z*v__2.z);
|
||||
__3.w = (v__1.w*v__2.w);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__2))[0] = __float22bfloat162_rn(((float2*)(&__3))[0]);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__2))[1] = __float22bfloat162_rn(((float2*)(&__3))[1]);
|
||||
*(uint2*)(output_local_cast + 0) = __2;
|
||||
*(uint2*)(output + (((((((((((int)threadIdx.x) >> 6) * 8192) + (((((int)threadIdx.x) & 15) >> 3) * 4096)) + (((int)blockIdx.y) * 1024)) + ((((int)threadIdx.x) & 7) * 128)) + (i_11 * 16)) + (((((int)threadIdx.x) & 63) >> 4) * 4)) + 67043328) - (((int)blockIdx.x) * 65536))) = *(uint2*)(output_local_cast + 0);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -0,0 +1,194 @@
|
|||
#include <tl_templates/maca/gemm.h>
|
||||
#include <tl_templates/maca/copy.h>
|
||||
#include <tl_templates/maca/reduce.h>
|
||||
#include <tl_templates/maca/intrin.h>
|
||||
#include <tl_templates/maca/atomic.h>
|
||||
#include <tl_templates/maca/threadblock_swizzle.h>
|
||||
#include <tl_templates/maca/debug.h>
|
||||
|
||||
extern "C" __global__ void packed_kernel_v_prefetch_after_qk_v15_kernel(const bfloat16_t* __restrict__ k, const int* __restrict__ kv_indptr, bfloat16_t* __restrict__ output, const bfloat16_t* __restrict__ q, const bfloat16_t* __restrict__ v);
|
||||
extern "C" __global__ void __launch_bounds__(512, 1) packed_kernel_v_prefetch_after_qk_v15_kernel(const bfloat16_t* __restrict__ k, const int* __restrict__ kv_indptr, bfloat16_t* __restrict__ output, const bfloat16_t* __restrict__ q, const bfloat16_t* __restrict__ v) {
|
||||
extern __shared__ __align__(1024) uchar buf_dyn_shmem[];
|
||||
float output_accum[32];
|
||||
float row_denom[1];
|
||||
float row_max[1];
|
||||
float row_scale[1];
|
||||
float scores[16];
|
||||
float row_max_prev[1];
|
||||
bfloat16_t probs[16];
|
||||
bfloat16_t output_local_cast[4];
|
||||
int copy_kv_len = kv_indptr[1];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + ((((((((((int)threadIdx.x) & 15) >> 3) * 8192) + (i * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8))) = *(uint4*)(q + ((((((i * 16384) + ((((int)threadIdx.x) >> 7) * 4096)) + (((int)blockIdx.y) * 1024)) + ((((int)threadIdx.x) & 127) * 8)) + 67043328) - (((int)blockIdx.x) * 65536)));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i_1 = 0; i_1 < 8; ++i_1) {
|
||||
float broadcast_var = 0x0p+0f/*0.000000e+00*/;
|
||||
*(float4*)(output_accum + (i_1 * 4)) = make_float4(broadcast_var, broadcast_var, broadcast_var, broadcast_var);
|
||||
}
|
||||
row_denom[0] = 0x0p+0f/*0.000000e+00*/;
|
||||
row_max[0] = -MACART_INF_F;
|
||||
for (int kv_tile = 0; kv_tile < ((16447 - (((int)blockIdx.x) * 16)) >> 6); ++kv_tile) {
|
||||
__syncthreads();
|
||||
if (((kv_tile * 64) + 64) <= copy_kv_len) {
|
||||
#pragma unroll
|
||||
for (int i_2 = 0; i_2 < 2; ++i_2) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_2 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(k + (((((kv_tile * 32768) + (i_2 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_3 = 0; i_3 < 2; ++i_3) {
|
||||
if ((((kv_tile * 64) + (i_3 * 32)) + (((int)threadIdx.x) >> 4)) < copy_kv_len) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_3 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(k + (((((kv_tile * 32768) + (i_3 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
} else {
|
||||
bfloat16_t broadcast_var_1 = bfloat16_t(0x0p+0f/*0.000000e+00*/);
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_3 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = make_uint4(__pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1));
|
||||
}
|
||||
}
|
||||
}
|
||||
if (kv_tile < ((16369 - (((int)blockIdx.x) * 16)) >> 6)) {
|
||||
#pragma unroll
|
||||
for (int i_4 = 0; i_4 < 4; ++i_4) {
|
||||
float broadcast_var_2 = 0x0p+0f/*0.000000e+00*/;
|
||||
*(float4*)(scores + (i_4 * 4)) = make_float4(broadcast_var_2, broadcast_var_2, broadcast_var_2, broadcast_var_2);
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_5 = 0; i_5 < 16; ++i_5) {
|
||||
float condval;
|
||||
if ((((((kv_tile * 64) + ((i_5 >> 2) * 16)) + (((((int)threadIdx.x) & 63) >> 4) * 4)) + (i_5 & 3)) <= (((((((int)threadIdx.x) >> 6) * 2) + ((((int)threadIdx.x) & 15) >> 3)) + 16368) - (((int)blockIdx.x) * 16)))) {
|
||||
condval = 0x0p+0f/*0.000000e+00*/;
|
||||
} else {
|
||||
condval = -MACART_INF_F;
|
||||
}
|
||||
scores[i_5] = condval;
|
||||
}
|
||||
}
|
||||
bfloat16_t A_local[4];
|
||||
bfloat16_t B_local[16];
|
||||
__syncthreads();
|
||||
for (int ki = 0; ki < 8; ++ki) {
|
||||
*(uint2*)(A_local + 0) = *(uint2*)(((bfloat16_t*)buf_dyn_shmem) + ((((((((ki >> 2) * 8192) + ((((int)threadIdx.x) >> 6) * 1024)) + ((((int)threadIdx.x) & 15) * 64)) + (((((((int)threadIdx.x) & 7) >> 2) + ((ki & 3) >> 1)) & 1) * 32)) + (((((((int)threadIdx.x) & 3) >> 1) + (ki & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 63) >> 5) + (((int)threadIdx.x) & 1)) & 1) * 8)) + (((((int)threadIdx.x) & 31) >> 4) * 4)));
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
*(uint2*)(B_local + (j * 4)) = *(uint2*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((ki >> 2) * 4096) + (j * 1024)) + ((((int)threadIdx.x) & 15) * 64)) + (((((((int)threadIdx.x) & 7) >> 2) + ((ki & 3) >> 1)) & 1) * 32)) + (((((((int)threadIdx.x) & 3) >> 1) + (ki & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 63) >> 5) + (((int)threadIdx.x) & 1)) & 1) * 8)) + (((((int)threadIdx.x) & 31) >> 4) * 4)) + 16384));
|
||||
}
|
||||
for (int j_1 = 0; j_1 < 4; ++j_1) {
|
||||
{
|
||||
*(((float32x4*)scores) + j_1) = __builtin_mxc_mma_16x16x16bf16(*(((bfloat16x4_vec*)B_local) + j_1),
|
||||
*(((bfloat16x4_vec*)A_local) + 0),
|
||||
*(((float32x4*)scores) + j_1));
|
||||
};
|
||||
}
|
||||
}
|
||||
row_max_prev[0] = -MACART_INF_F;
|
||||
#pragma unroll
|
||||
for (int rv = 0; rv < 16; ++rv) {
|
||||
row_max_prev[0] = max(row_max_prev[0], scores[(((rv & 3) * 4) + (rv >> 2))]);
|
||||
}
|
||||
__syncthreads();
|
||||
row_max_prev[0] = tl::AllReduce<tl::MaxOp, 64, 16, 0>::run(row_max_prev[0], (&(((float*)buf_dyn_shmem)[12800])));
|
||||
__syncthreads();
|
||||
if (((kv_tile * 64) + 64) <= copy_kv_len) {
|
||||
#pragma unroll
|
||||
for (int i_6 = 0; i_6 < 2; ++i_6) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_6 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(v + (((((kv_tile * 32768) + (i_6 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_7 = 0; i_7 < 2; ++i_7) {
|
||||
if ((((kv_tile * 64) + (i_7 * 32)) + (((int)threadIdx.x) >> 4)) < copy_kv_len) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_7 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(v + (((((kv_tile * 32768) + (i_7 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
} else {
|
||||
bfloat16_t broadcast_var_3 = bfloat16_t(0x0p+0f/*0.000000e+00*/);
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_7 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = make_uint4(__pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3));
|
||||
}
|
||||
}
|
||||
}
|
||||
row_scale[0] = exp2f(((row_max[0] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/) - (max(row_max[0], row_max_prev[0]) * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/)));
|
||||
row_max[0] = max(row_max[0], row_max_prev[0]);
|
||||
#pragma unroll
|
||||
for (int i_8 = 0; i_8 < 16; ++i_8) {
|
||||
scores[i_8] = exp2f(((scores[i_8] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/) - (row_max[0] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/)));
|
||||
}
|
||||
row_max_prev[0] = 0x0p+0f/*0.000000e+00*/;
|
||||
#pragma unroll
|
||||
for (int rv_1 = 0; rv_1 < 16; ++rv_1) {
|
||||
row_max_prev[0] = (row_max_prev[0] + scores[(((rv_1 & 3) * 4) + (rv_1 >> 2))]);
|
||||
}
|
||||
__syncthreads();
|
||||
row_max_prev[0] = tl::AllReduce<tl::SumOp, 64, 16, 0>::run(row_max_prev[0], (&(((float*)buf_dyn_shmem)[12288])));
|
||||
row_denom[0] = ((row_denom[0] * row_scale[0]) + row_max_prev[0]);
|
||||
#pragma unroll
|
||||
for (int i_9 = 0; i_9 < 4; ++i_9) {
|
||||
uint2 __1;
|
||||
float4 v_ = *(float4*)(scores + (i_9 * 4));
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__1))[0] = __float22bfloat162_rn(((float2*)(&v_))[0]);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__1))[1] = __float22bfloat162_rn(((float2*)(&v_))[1]);
|
||||
*(uint2*)(probs + (i_9 * 4)) = __1;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i_10 = 0; i_10 < 32; ++i_10) {
|
||||
output_accum[i_10] = (output_accum[i_10] * row_scale[0]);
|
||||
}
|
||||
bfloat16_t B_local_1[32];
|
||||
__syncthreads();
|
||||
for (int ki_1 = 0; ki_1 < 4; ++ki_1) {
|
||||
for (int j_2 = 0; j_2 < 8; ++j_2) {
|
||||
for (int local_id = 0; local_id < 4; ++local_id) {
|
||||
B_local_1[((j_2 * 4) + local_id)] = ((bfloat16_t*)buf_dyn_shmem)[((((((((((j_2 >> 2) * 4096) + (ki_1 * 1024)) + (((((int)threadIdx.x) & 63) >> 4) * 256)) + (local_id * 64)) + (((((((int)threadIdx.x) & 31) >> 4) + ((j_2 & 3) >> 1)) & 1) * 32)) + ((((local_id >> 1) + (j_2 & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 15) >> 3) + (local_id & 1)) & 1) * 8)) + (((int)threadIdx.x) & 7)) + 16384)];
|
||||
}
|
||||
}
|
||||
for (int j_3 = 0; j_3 < 8; ++j_3) {
|
||||
{
|
||||
*(((float32x4*)output_accum) + j_3) = __builtin_mxc_mma_16x16x16bf16(*(((bfloat16x4_vec*)B_local_1) + j_3),
|
||||
*(((bfloat16x4_vec*)probs) + ki_1),
|
||||
*(((float32x4*)output_accum) + j_3));
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
row_scale[0] = (0x1p+0f/*1.000000e+00*/ / row_denom[0]);
|
||||
#pragma unroll
|
||||
for (int i_11 = 0; i_11 < 8; ++i_11) {
|
||||
uint2 __2;
|
||||
float4 __3;
|
||||
float4 v__1 = *(float4*)(output_accum + (i_11 * 4));
|
||||
float4 v__2 = make_float4(row_scale[0], row_scale[0], row_scale[0], row_scale[0]);
|
||||
__3.x = (v__1.x*v__2.x);
|
||||
__3.y = (v__1.y*v__2.y);
|
||||
__3.z = (v__1.z*v__2.z);
|
||||
__3.w = (v__1.w*v__2.w);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__2))[0] = __float22bfloat162_rn(((float2*)(&__3))[0]);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__2))[1] = __float22bfloat162_rn(((float2*)(&__3))[1]);
|
||||
*(uint2*)(output_local_cast + 0) = __2;
|
||||
*(uint2*)(output + (((((((((((int)threadIdx.x) >> 6) * 8192) + (((((int)threadIdx.x) & 15) >> 3) * 4096)) + (((int)blockIdx.y) * 1024)) + ((((int)threadIdx.x) & 7) * 128)) + (i_11 * 16)) + (((((int)threadIdx.x) & 63) >> 4) * 4)) + 67043328) - (((int)blockIdx.x) * 65536))) = *(uint2*)(output_local_cast + 0);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
#define ERROR_BUF_SIZE 1024
|
||||
static char error_buf[ERROR_BUF_SIZE];
|
||||
|
||||
extern "C" const char* get_last_error() {
|
||||
return error_buf;
|
||||
}
|
||||
|
||||
extern "C" int init() {
|
||||
error_buf[0] = '\0';
|
||||
|
||||
if (53248 > 65536) {
|
||||
snprintf(error_buf, ERROR_BUF_SIZE, "Failed to set the allowed dynamic shared memory size for packed_kernel_v_prefetch_after_qk_v15_kernel to %d", 53248);
|
||||
return -1;
|
||||
}
|
||||
return 0;
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
extern "C" int call(bfloat16_t* __restrict__ q, bfloat16_t* __restrict__ k, bfloat16_t* __restrict__ v, bfloat16_t* __restrict__ output, int* __restrict__ qo_indptr, int* __restrict__ kv_indptr, mcStream_t stream=mcStreamDefault) {
|
||||
packed_kernel_v_prefetch_after_qk_v15_kernel<<<dim3(1024, 4, 1), dim3(512, 1, 1), 53248, stream>>>(k, kv_indptr, output, q, v);
|
||||
TILELANG_CHECK_LAST_ERROR("packed_kernel_v_prefetch_after_qk_v15_kernel");
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -0,0 +1,167 @@
|
|||
#include <tl_templates/maca/gemm.h>
|
||||
#include <tl_templates/maca/copy.h>
|
||||
#include <tl_templates/maca/reduce.h>
|
||||
#include <tl_templates/maca/intrin.h>
|
||||
#include <tl_templates/maca/atomic.h>
|
||||
#include <tl_templates/maca/threadblock_swizzle.h>
|
||||
#include <tl_templates/maca/debug.h>
|
||||
|
||||
extern "C" __global__ void packed_kernel_factor_softmax_scale_v18_kernel(const bfloat16_t* __restrict__ k, const int* __restrict__ kv_indptr, bfloat16_t* __restrict__ output, const bfloat16_t* __restrict__ q, const bfloat16_t* __restrict__ v);
|
||||
extern "C" __global__ void __launch_bounds__(512, 1) packed_kernel_factor_softmax_scale_v18_kernel(const bfloat16_t* __restrict__ k, const int* __restrict__ kv_indptr, bfloat16_t* __restrict__ output, const bfloat16_t* __restrict__ q, const bfloat16_t* __restrict__ v) {
|
||||
extern __shared__ __align__(1024) uchar buf_dyn_shmem[];
|
||||
float output_accum[32];
|
||||
float row_denom[1];
|
||||
float row_max[1];
|
||||
float row_scale[1];
|
||||
float scores[16];
|
||||
float row_max_prev[1];
|
||||
bfloat16_t probs[16];
|
||||
bfloat16_t output_local_cast[4];
|
||||
int copy_kv_len = kv_indptr[1];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + ((((((((((int)threadIdx.x) & 15) >> 3) * 8192) + (i * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8))) = *(uint4*)(q + ((((((i * 16384) + ((((int)threadIdx.x) >> 7) * 4096)) + (((int)blockIdx.y) * 1024)) + ((((int)threadIdx.x) & 127) * 8)) + 67043328) - (((int)blockIdx.x) * 65536)));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i_1 = 0; i_1 < 8; ++i_1) {
|
||||
float broadcast_var = 0x0p+0f/*0.000000e+00*/;
|
||||
*(float4*)(output_accum + (i_1 * 4)) = make_float4(broadcast_var, broadcast_var, broadcast_var, broadcast_var);
|
||||
}
|
||||
row_denom[0] = 0x0p+0f/*0.000000e+00*/;
|
||||
row_max[0] = -MACART_INF_F;
|
||||
for (int kv_tile = 0; kv_tile < ((16447 - (((int)blockIdx.x) * 16)) >> 6); ++kv_tile) {
|
||||
__syncthreads();
|
||||
if (((kv_tile * 64) + 64) <= copy_kv_len) {
|
||||
#pragma unroll
|
||||
for (int i_2 = 0; i_2 < 2; ++i_2) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_2 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(k + (((((kv_tile * 32768) + (i_2 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_3 = 0; i_3 < 2; ++i_3) {
|
||||
if ((((kv_tile * 64) + (i_3 * 32)) + (((int)threadIdx.x) >> 4)) < copy_kv_len) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_3 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(k + (((((kv_tile * 32768) + (i_3 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
} else {
|
||||
bfloat16_t broadcast_var_1 = bfloat16_t(0x0p+0f/*0.000000e+00*/);
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_3 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = make_uint4(__pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1));
|
||||
}
|
||||
}
|
||||
}
|
||||
if (kv_tile < ((16369 - (((int)blockIdx.x) * 16)) >> 6)) {
|
||||
#pragma unroll
|
||||
for (int i_4 = 0; i_4 < 4; ++i_4) {
|
||||
float broadcast_var_2 = 0x0p+0f/*0.000000e+00*/;
|
||||
*(float4*)(scores + (i_4 * 4)) = make_float4(broadcast_var_2, broadcast_var_2, broadcast_var_2, broadcast_var_2);
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_5 = 0; i_5 < 16; ++i_5) {
|
||||
float condval;
|
||||
if ((((((kv_tile * 64) + ((i_5 >> 2) * 16)) + (((((int)threadIdx.x) & 63) >> 4) * 4)) + (i_5 & 3)) <= (((((((int)threadIdx.x) >> 6) * 2) + ((((int)threadIdx.x) & 15) >> 3)) + 16368) - (((int)blockIdx.x) * 16)))) {
|
||||
condval = 0x0p+0f/*0.000000e+00*/;
|
||||
} else {
|
||||
condval = -MACART_INF_F;
|
||||
}
|
||||
scores[i_5] = condval;
|
||||
}
|
||||
}
|
||||
bfloat16_t A_local[4];
|
||||
bfloat16_t B_local[16];
|
||||
__syncthreads();
|
||||
for (int ki = 0; ki < 8; ++ki) {
|
||||
*(uint2*)(A_local + 0) = *(uint2*)(((bfloat16_t*)buf_dyn_shmem) + ((((((((ki >> 2) * 8192) + ((((int)threadIdx.x) >> 6) * 1024)) + ((((int)threadIdx.x) & 15) * 64)) + (((((((int)threadIdx.x) & 7) >> 2) + ((ki & 3) >> 1)) & 1) * 32)) + (((((((int)threadIdx.x) & 3) >> 1) + (ki & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 63) >> 5) + (((int)threadIdx.x) & 1)) & 1) * 8)) + (((((int)threadIdx.x) & 31) >> 4) * 4)));
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
*(uint2*)(B_local + (j * 4)) = *(uint2*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((ki >> 2) * 4096) + (j * 1024)) + ((((int)threadIdx.x) & 15) * 64)) + (((((((int)threadIdx.x) & 7) >> 2) + ((ki & 3) >> 1)) & 1) * 32)) + (((((((int)threadIdx.x) & 3) >> 1) + (ki & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 63) >> 5) + (((int)threadIdx.x) & 1)) & 1) * 8)) + (((((int)threadIdx.x) & 31) >> 4) * 4)) + 16384));
|
||||
}
|
||||
for (int j_1 = 0; j_1 < 4; ++j_1) {
|
||||
{
|
||||
*(((float32x4*)scores) + j_1) = __builtin_mxc_mma_16x16x16bf16(*(((bfloat16x4_vec*)B_local) + j_1),
|
||||
*(((bfloat16x4_vec*)A_local) + 0),
|
||||
*(((float32x4*)scores) + j_1));
|
||||
};
|
||||
}
|
||||
}
|
||||
row_max_prev[0] = -MACART_INF_F;
|
||||
#pragma unroll
|
||||
for (int rv = 0; rv < 16; ++rv) {
|
||||
row_max_prev[0] = max(row_max_prev[0], scores[(((rv & 3) * 4) + (rv >> 2))]);
|
||||
}
|
||||
__syncthreads();
|
||||
row_max_prev[0] = tl::AllReduce<tl::MaxOp, 64, 16, 0>::run(row_max_prev[0], (&(((float*)buf_dyn_shmem)[12800])));
|
||||
row_scale[0] = exp2f(((row_max[0] - max(row_max[0], row_max_prev[0])) * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/));
|
||||
row_max[0] = max(row_max[0], row_max_prev[0]);
|
||||
#pragma unroll
|
||||
for (int i_6 = 0; i_6 < 16; ++i_6) {
|
||||
scores[i_6] = exp2f(((scores[i_6] - row_max[0]) * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/));
|
||||
}
|
||||
row_max_prev[0] = 0x0p+0f/*0.000000e+00*/;
|
||||
#pragma unroll
|
||||
for (int rv_1 = 0; rv_1 < 16; ++rv_1) {
|
||||
row_max_prev[0] = (row_max_prev[0] + scores[(((rv_1 & 3) * 4) + (rv_1 >> 2))]);
|
||||
}
|
||||
row_max_prev[0] = tl::AllReduce<tl::SumOp, 64, 16, 0>::run(row_max_prev[0], (&(((float*)buf_dyn_shmem)[12288])));
|
||||
row_denom[0] = ((row_denom[0] * row_scale[0]) + row_max_prev[0]);
|
||||
#pragma unroll
|
||||
for (int i_7 = 0; i_7 < 4; ++i_7) {
|
||||
uint2 __1;
|
||||
float4 v_ = *(float4*)(scores + (i_7 * 4));
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__1))[0] = __float22bfloat162_rn(((float2*)(&v_))[0]);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__1))[1] = __float22bfloat162_rn(((float2*)(&v_))[1]);
|
||||
*(uint2*)(probs + (i_7 * 4)) = __1;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i_8 = 0; i_8 < 32; ++i_8) {
|
||||
output_accum[i_8] = (output_accum[i_8] * row_scale[0]);
|
||||
}
|
||||
__syncthreads();
|
||||
if (((kv_tile * 64) + 64) <= copy_kv_len) {
|
||||
#pragma unroll
|
||||
for (int i_9 = 0; i_9 < 2; ++i_9) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_9 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(v + (((((kv_tile * 32768) + (i_9 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_10 = 0; i_10 < 2; ++i_10) {
|
||||
if ((((kv_tile * 64) + (i_10 * 32)) + (((int)threadIdx.x) >> 4)) < copy_kv_len) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_10 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(v + (((((kv_tile * 32768) + (i_10 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
} else {
|
||||
bfloat16_t broadcast_var_3 = bfloat16_t(0x0p+0f/*0.000000e+00*/);
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_10 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = make_uint4(__pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3));
|
||||
}
|
||||
}
|
||||
}
|
||||
bfloat16_t B_local_1[32];
|
||||
__syncthreads();
|
||||
for (int ki_1 = 0; ki_1 < 4; ++ki_1) {
|
||||
for (int j_2 = 0; j_2 < 8; ++j_2) {
|
||||
for (int local_id = 0; local_id < 4; ++local_id) {
|
||||
B_local_1[((j_2 * 4) + local_id)] = ((bfloat16_t*)buf_dyn_shmem)[((((((((((j_2 >> 2) * 4096) + (ki_1 * 1024)) + (((((int)threadIdx.x) & 63) >> 4) * 256)) + (local_id * 64)) + (((((((int)threadIdx.x) & 31) >> 4) + ((j_2 & 3) >> 1)) & 1) * 32)) + ((((local_id >> 1) + (j_2 & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 15) >> 3) + (local_id & 1)) & 1) * 8)) + (((int)threadIdx.x) & 7)) + 16384)];
|
||||
}
|
||||
}
|
||||
for (int j_3 = 0; j_3 < 8; ++j_3) {
|
||||
{
|
||||
*(((float32x4*)output_accum) + j_3) = __builtin_mxc_mma_16x16x16bf16(*(((bfloat16x4_vec*)B_local_1) + j_3),
|
||||
*(((bfloat16x4_vec*)probs) + ki_1),
|
||||
*(((float32x4*)output_accum) + j_3));
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
row_scale[0] = (0x1p+0f/*1.000000e+00*/ / row_denom[0]);
|
||||
#pragma unroll
|
||||
for (int i_11 = 0; i_11 < 8; ++i_11) {
|
||||
uint2 __2;
|
||||
float4 __3;
|
||||
float4 v__1 = *(float4*)(output_accum + (i_11 * 4));
|
||||
float4 v__2 = make_float4(row_scale[0], row_scale[0], row_scale[0], row_scale[0]);
|
||||
__3.x = (v__1.x*v__2.x);
|
||||
__3.y = (v__1.y*v__2.y);
|
||||
__3.z = (v__1.z*v__2.z);
|
||||
__3.w = (v__1.w*v__2.w);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__2))[0] = __float22bfloat162_rn(((float2*)(&__3))[0]);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__2))[1] = __float22bfloat162_rn(((float2*)(&__3))[1]);
|
||||
*(uint2*)(output_local_cast + 0) = __2;
|
||||
*(uint2*)(output + (((((((((((int)threadIdx.x) >> 6) * 8192) + (((((int)threadIdx.x) & 15) >> 3) * 4096)) + (((int)blockIdx.y) * 1024)) + ((((int)threadIdx.x) & 7) * 128)) + (i_11 * 16)) + (((((int)threadIdx.x) & 63) >> 4) * 4)) + 67043328) - (((int)blockIdx.x) * 65536))) = *(uint2*)(output_local_cast + 0);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -0,0 +1,193 @@
|
|||
#include <tl_templates/maca/gemm.h>
|
||||
#include <tl_templates/maca/copy.h>
|
||||
#include <tl_templates/maca/reduce.h>
|
||||
#include <tl_templates/maca/intrin.h>
|
||||
#include <tl_templates/maca/atomic.h>
|
||||
#include <tl_templates/maca/threadblock_swizzle.h>
|
||||
#include <tl_templates/maca/debug.h>
|
||||
|
||||
extern "C" __global__ void packed_kernel_factor_softmax_scale_v18_kernel(const bfloat16_t* __restrict__ k, const int* __restrict__ kv_indptr, bfloat16_t* __restrict__ output, const bfloat16_t* __restrict__ q, const bfloat16_t* __restrict__ v);
|
||||
extern "C" __global__ void __launch_bounds__(512, 1) packed_kernel_factor_softmax_scale_v18_kernel(const bfloat16_t* __restrict__ k, const int* __restrict__ kv_indptr, bfloat16_t* __restrict__ output, const bfloat16_t* __restrict__ q, const bfloat16_t* __restrict__ v) {
|
||||
extern __shared__ __align__(1024) uchar buf_dyn_shmem[];
|
||||
float output_accum[32];
|
||||
float row_denom[1];
|
||||
float row_max[1];
|
||||
float row_scale[1];
|
||||
float scores[16];
|
||||
float row_max_prev[1];
|
||||
bfloat16_t probs[16];
|
||||
bfloat16_t output_local_cast[4];
|
||||
int copy_kv_len = kv_indptr[1];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + ((((((((((int)threadIdx.x) & 15) >> 3) * 8192) + (i * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8))) = *(uint4*)(q + ((((((i * 16384) + ((((int)threadIdx.x) >> 7) * 4096)) + (((int)blockIdx.y) * 1024)) + ((((int)threadIdx.x) & 127) * 8)) + 67043328) - (((int)blockIdx.x) * 65536)));
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i_1 = 0; i_1 < 8; ++i_1) {
|
||||
float broadcast_var = 0x0p+0f/*0.000000e+00*/;
|
||||
*(float4*)(output_accum + (i_1 * 4)) = make_float4(broadcast_var, broadcast_var, broadcast_var, broadcast_var);
|
||||
}
|
||||
row_denom[0] = 0x0p+0f/*0.000000e+00*/;
|
||||
row_max[0] = -MACART_INF_F;
|
||||
for (int kv_tile = 0; kv_tile < ((16447 - (((int)blockIdx.x) * 16)) >> 6); ++kv_tile) {
|
||||
__syncthreads();
|
||||
if (((kv_tile * 64) + 64) <= copy_kv_len) {
|
||||
#pragma unroll
|
||||
for (int i_2 = 0; i_2 < 2; ++i_2) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_2 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(k + (((((kv_tile * 32768) + (i_2 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_3 = 0; i_3 < 2; ++i_3) {
|
||||
if ((((kv_tile * 64) + (i_3 * 32)) + (((int)threadIdx.x) >> 4)) < copy_kv_len) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_3 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(k + (((((kv_tile * 32768) + (i_3 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
} else {
|
||||
bfloat16_t broadcast_var_1 = bfloat16_t(0x0p+0f/*0.000000e+00*/);
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_3 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = make_uint4(__pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1), __pack_maca_bfloat162(broadcast_var_1, broadcast_var_1));
|
||||
}
|
||||
}
|
||||
}
|
||||
if (kv_tile < ((16369 - (((int)blockIdx.x) * 16)) >> 6)) {
|
||||
#pragma unroll
|
||||
for (int i_4 = 0; i_4 < 4; ++i_4) {
|
||||
float broadcast_var_2 = 0x0p+0f/*0.000000e+00*/;
|
||||
*(float4*)(scores + (i_4 * 4)) = make_float4(broadcast_var_2, broadcast_var_2, broadcast_var_2, broadcast_var_2);
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_5 = 0; i_5 < 16; ++i_5) {
|
||||
float condval;
|
||||
if ((((((kv_tile * 64) + ((i_5 >> 2) * 16)) + (((((int)threadIdx.x) & 63) >> 4) * 4)) + (i_5 & 3)) <= (((((((int)threadIdx.x) >> 6) * 2) + ((((int)threadIdx.x) & 15) >> 3)) + 16368) - (((int)blockIdx.x) * 16)))) {
|
||||
condval = 0x0p+0f/*0.000000e+00*/;
|
||||
} else {
|
||||
condval = -MACART_INF_F;
|
||||
}
|
||||
scores[i_5] = condval;
|
||||
}
|
||||
}
|
||||
bfloat16_t A_local[4];
|
||||
bfloat16_t B_local[16];
|
||||
__syncthreads();
|
||||
for (int ki = 0; ki < 8; ++ki) {
|
||||
*(uint2*)(A_local + 0) = *(uint2*)(((bfloat16_t*)buf_dyn_shmem) + ((((((((ki >> 2) * 8192) + ((((int)threadIdx.x) >> 6) * 1024)) + ((((int)threadIdx.x) & 15) * 64)) + (((((((int)threadIdx.x) & 7) >> 2) + ((ki & 3) >> 1)) & 1) * 32)) + (((((((int)threadIdx.x) & 3) >> 1) + (ki & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 63) >> 5) + (((int)threadIdx.x) & 1)) & 1) * 8)) + (((((int)threadIdx.x) & 31) >> 4) * 4)));
|
||||
for (int j = 0; j < 4; ++j) {
|
||||
*(uint2*)(B_local + (j * 4)) = *(uint2*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((ki >> 2) * 4096) + (j * 1024)) + ((((int)threadIdx.x) & 15) * 64)) + (((((((int)threadIdx.x) & 7) >> 2) + ((ki & 3) >> 1)) & 1) * 32)) + (((((((int)threadIdx.x) & 3) >> 1) + (ki & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 63) >> 5) + (((int)threadIdx.x) & 1)) & 1) * 8)) + (((((int)threadIdx.x) & 31) >> 4) * 4)) + 16384));
|
||||
}
|
||||
for (int j_1 = 0; j_1 < 4; ++j_1) {
|
||||
{
|
||||
*(((float32x4*)scores) + j_1) = __builtin_mxc_mma_16x16x16bf16(*(((bfloat16x4_vec*)B_local) + j_1),
|
||||
*(((bfloat16x4_vec*)A_local) + 0),
|
||||
*(((float32x4*)scores) + j_1));
|
||||
};
|
||||
}
|
||||
}
|
||||
row_max_prev[0] = -MACART_INF_F;
|
||||
#pragma unroll
|
||||
for (int rv = 0; rv < 16; ++rv) {
|
||||
row_max_prev[0] = max(row_max_prev[0], scores[(((rv & 3) * 4) + (rv >> 2))]);
|
||||
}
|
||||
__syncthreads();
|
||||
row_max_prev[0] = tl::AllReduce<tl::MaxOp, 64, 16, 0>::run(row_max_prev[0], (&(((float*)buf_dyn_shmem)[12800])));
|
||||
row_scale[0] = exp2f(((row_max[0] - max(row_max[0], row_max_prev[0])) * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/));
|
||||
row_max[0] = max(row_max[0], row_max_prev[0]);
|
||||
#pragma unroll
|
||||
for (int i_6 = 0; i_6 < 16; ++i_6) {
|
||||
scores[i_6] = exp2f(((scores[i_6] - row_max[0]) * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/));
|
||||
}
|
||||
row_max_prev[0] = 0x0p+0f/*0.000000e+00*/;
|
||||
#pragma unroll
|
||||
for (int rv_1 = 0; rv_1 < 16; ++rv_1) {
|
||||
row_max_prev[0] = (row_max_prev[0] + scores[(((rv_1 & 3) * 4) + (rv_1 >> 2))]);
|
||||
}
|
||||
row_max_prev[0] = tl::AllReduce<tl::SumOp, 64, 16, 0>::run(row_max_prev[0], (&(((float*)buf_dyn_shmem)[12288])));
|
||||
row_denom[0] = ((row_denom[0] * row_scale[0]) + row_max_prev[0]);
|
||||
#pragma unroll
|
||||
for (int i_7 = 0; i_7 < 4; ++i_7) {
|
||||
uint2 __1;
|
||||
float4 v_ = *(float4*)(scores + (i_7 * 4));
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__1))[0] = __float22bfloat162_rn(((float2*)(&v_))[0]);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__1))[1] = __float22bfloat162_rn(((float2*)(&v_))[1]);
|
||||
*(uint2*)(probs + (i_7 * 4)) = __1;
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i_8 = 0; i_8 < 32; ++i_8) {
|
||||
output_accum[i_8] = (output_accum[i_8] * row_scale[0]);
|
||||
}
|
||||
__syncthreads();
|
||||
if (((kv_tile * 64) + 64) <= copy_kv_len) {
|
||||
#pragma unroll
|
||||
for (int i_9 = 0; i_9 < 2; ++i_9) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_9 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(v + (((((kv_tile * 32768) + (i_9 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
}
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (int i_10 = 0; i_10 < 2; ++i_10) {
|
||||
if ((((kv_tile * 64) + (i_10 * 32)) + (((int)threadIdx.x) >> 4)) < copy_kv_len) {
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_10 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = *(uint4*)(v + (((((kv_tile * 32768) + (i_10 * 16384)) + ((((int)threadIdx.x) >> 4) * 512)) + (((int)blockIdx.y) * 128)) + ((((int)threadIdx.x) & 15) * 8)));
|
||||
} else {
|
||||
bfloat16_t broadcast_var_3 = bfloat16_t(0x0p+0f/*0.000000e+00*/);
|
||||
*(uint4*)(((bfloat16_t*)buf_dyn_shmem) + (((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_10 * 2048)) + ((((int)threadIdx.x) >> 4) * 64)) + (((((((int)threadIdx.x) & 127) >> 6) + ((((int)threadIdx.x) & 7) >> 2)) & 1) * 32)) + (((((((int)threadIdx.x) & 63) >> 5) + ((((int)threadIdx.x) & 3) >> 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8)) + 16384)) = make_uint4(__pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3), __pack_maca_bfloat162(broadcast_var_3, broadcast_var_3));
|
||||
}
|
||||
}
|
||||
}
|
||||
bfloat16_t B_local_1[32];
|
||||
__syncthreads();
|
||||
for (int ki_1 = 0; ki_1 < 4; ++ki_1) {
|
||||
for (int j_2 = 0; j_2 < 8; ++j_2) {
|
||||
for (int local_id = 0; local_id < 4; ++local_id) {
|
||||
B_local_1[((j_2 * 4) + local_id)] = ((bfloat16_t*)buf_dyn_shmem)[((((((((((j_2 >> 2) * 4096) + (ki_1 * 1024)) + (((((int)threadIdx.x) & 63) >> 4) * 256)) + (local_id * 64)) + (((((((int)threadIdx.x) & 31) >> 4) + ((j_2 & 3) >> 1)) & 1) * 32)) + ((((local_id >> 1) + (j_2 & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 15) >> 3) + (local_id & 1)) & 1) * 8)) + (((int)threadIdx.x) & 7)) + 16384)];
|
||||
}
|
||||
}
|
||||
for (int j_3 = 0; j_3 < 8; ++j_3) {
|
||||
{
|
||||
*(((float32x4*)output_accum) + j_3) = __builtin_mxc_mma_16x16x16bf16(*(((bfloat16x4_vec*)B_local_1) + j_3),
|
||||
*(((bfloat16x4_vec*)probs) + ki_1),
|
||||
*(((float32x4*)output_accum) + j_3));
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
row_scale[0] = (0x1p+0f/*1.000000e+00*/ / row_denom[0]);
|
||||
#pragma unroll
|
||||
for (int i_11 = 0; i_11 < 8; ++i_11) {
|
||||
uint2 __2;
|
||||
float4 __3;
|
||||
float4 v__1 = *(float4*)(output_accum + (i_11 * 4));
|
||||
float4 v__2 = make_float4(row_scale[0], row_scale[0], row_scale[0], row_scale[0]);
|
||||
__3.x = (v__1.x*v__2.x);
|
||||
__3.y = (v__1.y*v__2.y);
|
||||
__3.z = (v__1.z*v__2.z);
|
||||
__3.w = (v__1.w*v__2.w);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__2))[0] = __float22bfloat162_rn(((float2*)(&__3))[0]);
|
||||
(reinterpret_cast<__maca_bfloat162*>(&__2))[1] = __float22bfloat162_rn(((float2*)(&__3))[1]);
|
||||
*(uint2*)(output_local_cast + 0) = __2;
|
||||
*(uint2*)(output + (((((((((((int)threadIdx.x) >> 6) * 8192) + (((((int)threadIdx.x) & 15) >> 3) * 4096)) + (((int)blockIdx.y) * 1024)) + ((((int)threadIdx.x) & 7) * 128)) + (i_11 * 16)) + (((((int)threadIdx.x) & 63) >> 4) * 4)) + 67043328) - (((int)blockIdx.x) * 65536))) = *(uint2*)(output_local_cast + 0);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
#define ERROR_BUF_SIZE 1024
|
||||
static char error_buf[ERROR_BUF_SIZE];
|
||||
|
||||
extern "C" const char* get_last_error() {
|
||||
return error_buf;
|
||||
}
|
||||
|
||||
extern "C" int init() {
|
||||
error_buf[0] = '\0';
|
||||
|
||||
if (53248 > 65536) {
|
||||
snprintf(error_buf, ERROR_BUF_SIZE, "Failed to set the allowed dynamic shared memory size for packed_kernel_factor_softmax_scale_v18_kernel to %d", 53248);
|
||||
return -1;
|
||||
}
|
||||
return 0;
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
extern "C" int call(bfloat16_t* __restrict__ q, bfloat16_t* __restrict__ k, bfloat16_t* __restrict__ v, bfloat16_t* __restrict__ output, int* __restrict__ qo_indptr, int* __restrict__ kv_indptr, mcStream_t stream=mcStreamDefault) {
|
||||
packed_kernel_factor_softmax_scale_v18_kernel<<<dim3(1024, 4, 1), dim3(512, 1, 1), 53248, stream>>>(k, kv_indptr, output, q, v);
|
||||
TILELANG_CHECK_LAST_ERROR("packed_kernel_factor_softmax_scale_v18_kernel");
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -0,0 +1,178 @@
|
|||
import csv
|
||||
import importlib.util
|
||||
import statistics
|
||||
import time
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
KERNEL_PATH = "/data/operator_task_package/flashinfer_task_package/kernel_ops/FlashinferRaggedPrefill/tilelang/opt_012_dense_softmax_cleanup.py"
|
||||
OUTPUT_PATH = "/data/operator_task_package/flashinfer_task_package/kernel_ops/FlashinferRaggedPrefill/results/tilelang_64g/opt_012_benchmark.csv"
|
||||
SAMPLE_COUNT = 15
|
||||
TARGET_SAMPLE_MS = 100.0
|
||||
MAX_REPEATS = 2000
|
||||
WARMUP_REPEATS = 10
|
||||
|
||||
CASES = (
|
||||
(1, "ragged_b33_total16294", [987] + [478] * 31 + [489], [987] + [478] * 31 + [489]),
|
||||
(2, "equal_b1_s1024", [1024], [1024]),
|
||||
(3, "equal_b1_s4096", [4096], [4096]),
|
||||
(4, "equal_b1_s16384", [16384], [16384]),
|
||||
(5, "equal_b4_s1024", [1024] * 4, [1024] * 4),
|
||||
(6, "equal_b4_s4096", [4096] * 4, [4096] * 4),
|
||||
(7, "equal_b16_s1024", [1024] * 16, [1024] * 16),
|
||||
(8, "equal_b16_s2048", [2048] * 16, [2048] * 16),
|
||||
(9, "varlen_uniform_q512_k1024_b4", [512] * 4, [1024] * 4),
|
||||
(10, "varlen_mixed_b4", [640, 384, 256, 256], [1280, 1024, 768, 512]),
|
||||
(11, "varlen_q_lt_kv_b2", [512, 512], [2048, 1024]),
|
||||
(12, "ragged_b27_total12251", [873] + [438] * 25 + [428], [873] + [438] * 25 + [428]),
|
||||
(13, "short_ragged_b15_total969", [123] + [60] * 13 + [66], [123] + [60] * 13 + [66]),
|
||||
(14, "single_token", [1], [1]),
|
||||
(15, "tail_non_power2", [65, 33], [65, 33]),
|
||||
)
|
||||
|
||||
|
||||
def load_kernel():
|
||||
spec = importlib.util.spec_from_file_location("ragged_tilelang_opt012", KERNEL_PATH)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def make_indptr(lengths):
|
||||
values = [0]
|
||||
for length in lengths:
|
||||
values.append(values[-1] + length)
|
||||
return torch.tensor(values, device="cuda", dtype=torch.int32)
|
||||
|
||||
|
||||
def event_average_ms(function, repeats):
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(repeats):
|
||||
function()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) / repeats
|
||||
|
||||
|
||||
def percentile(values, fraction):
|
||||
ordered = sorted(values)
|
||||
position = fraction * (len(ordered) - 1)
|
||||
lower = int(position)
|
||||
upper = min(lower + 1, len(ordered) - 1)
|
||||
weight = position - lower
|
||||
return ordered[lower] * (1.0 - weight) + ordered[upper] * weight
|
||||
|
||||
|
||||
def benchmark_case(module, case_id, name, q_lengths, kv_lengths):
|
||||
torch.manual_seed(20260715 + case_id)
|
||||
batch_size = len(q_lengths)
|
||||
total_q = sum(q_lengths)
|
||||
total_kv = sum(kv_lengths)
|
||||
max_q = max(q_lengths)
|
||||
max_kv = max(kv_lengths)
|
||||
seq_len = max(max_q, max_kv)
|
||||
|
||||
q = torch.randn(total_q, 32, 128, device="cuda", dtype=torch.bfloat16)
|
||||
k = torch.randn(total_kv, 4, 128, device="cuda", dtype=torch.bfloat16)
|
||||
v = torch.randn_like(k)
|
||||
output = torch.empty_like(q)
|
||||
qo_indptr = make_indptr(q_lengths)
|
||||
kv_indptr = make_indptr(kv_lengths)
|
||||
|
||||
def launch():
|
||||
module.run_kernel(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
batch_size,
|
||||
seq_len,
|
||||
32,
|
||||
4,
|
||||
128,
|
||||
128,
|
||||
1,
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
compile_start = time.perf_counter()
|
||||
launch()
|
||||
torch.cuda.synchronize()
|
||||
compile_seconds = time.perf_counter() - compile_start
|
||||
|
||||
for _ in range(WARMUP_REPEATS):
|
||||
launch()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
pilot_ms = event_average_ms(launch, 3)
|
||||
repeats = max(1, min(MAX_REPEATS, int(TARGET_SAMPLE_MS / max(pilot_ms, 1.0e-4))))
|
||||
samples = [event_average_ms(launch, repeats) for _ in range(SAMPLE_COUNT)]
|
||||
|
||||
median_ms = statistics.median(samples)
|
||||
p10_ms = percentile(samples, 0.10)
|
||||
p90_ms = percentile(samples, 0.90)
|
||||
spread_pct = (p90_ms - p10_ms) / median_ms * 100.0
|
||||
|
||||
row = {
|
||||
"case_id": case_id,
|
||||
"config": name,
|
||||
"batch": batch_size,
|
||||
"total_q": total_q,
|
||||
"total_kv": total_kv,
|
||||
"max_q": max_q,
|
||||
"max_kv": max_kv,
|
||||
"seq_len": seq_len,
|
||||
"compile_first_launch_s": f"{compile_seconds:.6f}",
|
||||
"pilot_ms": f"{pilot_ms:.6f}",
|
||||
"repeats_per_sample": repeats,
|
||||
"sample_count": SAMPLE_COUNT,
|
||||
"median_ms": f"{median_ms:.6f}",
|
||||
"p10_ms": f"{p10_ms:.6f}",
|
||||
"p90_ms": f"{p90_ms:.6f}",
|
||||
"min_ms": f"{min(samples):.6f}",
|
||||
"max_ms": f"{max(samples):.6f}",
|
||||
"spread_pct": f"{spread_pct:.3f}",
|
||||
}
|
||||
|
||||
print(
|
||||
f"case={case_id:02d} {name}: median={median_ms:.6f} ms, "
|
||||
f"p10={p10_ms:.6f}, p90={p90_ms:.6f}, spread={spread_pct:.2f}%, "
|
||||
f"repeats={repeats}, compile_first={compile_seconds:.3f} s",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
del q, k, v, output, qo_indptr, kv_indptr
|
||||
torch.cuda.empty_cache()
|
||||
return row
|
||||
|
||||
|
||||
def main():
|
||||
module = load_kernel()
|
||||
rows = []
|
||||
fieldnames = None
|
||||
|
||||
for case in CASES:
|
||||
row = benchmark_case(module, *case)
|
||||
rows.append(row)
|
||||
if fieldnames is None:
|
||||
fieldnames = list(row.keys())
|
||||
with open(OUTPUT_PATH, "w", newline="", encoding="ascii") as output_file:
|
||||
writer = csv.DictWriter(output_file, fieldnames=fieldnames)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
total_median_ms = sum(float(row["median_ms"]) for row in rows)
|
||||
print(
|
||||
f"TILELANG_OPT012_BENCHMARK_PASS cases={len(rows)} "
|
||||
f"sum_of_medians_ms={total_median_ms:.6f} output={OUTPUT_PATH}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,108 @@
|
|||
import importlib.util
|
||||
import math
|
||||
import torch
|
||||
|
||||
ATOL = 1.6e-2
|
||||
RTOL = 1.6e-2
|
||||
|
||||
CASES = (
|
||||
(1, [987] + [478] * 31 + [489], [987] + [478] * 31 + [489]),
|
||||
(2, [1024], [1024]),
|
||||
(3, [4096], [4096]),
|
||||
(4, [16384], [16384]),
|
||||
(5, [1024] * 4, [1024] * 4),
|
||||
(6, [4096] * 4, [4096] * 4),
|
||||
(7, [1024] * 16, [1024] * 16),
|
||||
(8, [2048] * 16, [2048] * 16),
|
||||
(9, [512] * 4, [1024] * 4),
|
||||
(10, [640, 384, 256, 256], [1280, 1024, 768, 512]),
|
||||
(11, [512, 512], [2048, 1024]),
|
||||
(12, [873] + [438] * 25 + [428], [873] + [438] * 25 + [428]),
|
||||
(13, [123] + [60] * 13 + [66], [123] + [60] * 13 + [66]),
|
||||
(14, [1], [1]),
|
||||
(15, [65, 33], [65, 33]),
|
||||
)
|
||||
|
||||
|
||||
def load_kernel(kernel_path):
|
||||
spec = importlib.util.spec_from_file_location("ragged_kernel", kernel_path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def indptr(lengths):
|
||||
values = [0]
|
||||
for length in lengths:
|
||||
values.append(values[-1] + length)
|
||||
return torch.tensor(values, dtype=torch.int32, device="cuda")
|
||||
|
||||
|
||||
def main():
|
||||
import argparse
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--kernel-path", required=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
torch.manual_seed(20260714)
|
||||
module = load_kernel(args.kernel_path)
|
||||
|
||||
print("case,batch,total_q,total_kv,max_len,match,max_abs,worst_ratio,nan,inf,pass")
|
||||
for case_id, q_lens, kv_lens in CASES:
|
||||
batch_size = len(q_lens)
|
||||
total_q = sum(q_lens)
|
||||
total_kv = sum(kv_lens)
|
||||
max_len = max(max(q_lens), max(kv_lens))
|
||||
qo = indptr(q_lens)
|
||||
kv = indptr(kv_lens)
|
||||
q = torch.randn(total_q, 32, 128, dtype=torch.bfloat16, device="cuda")
|
||||
k = torch.randn(total_kv, 4, 128, dtype=torch.bfloat16, device="cuda")
|
||||
v = torch.randn(total_kv, 4, 128, dtype=torch.bfloat16, device="cuda")
|
||||
output = torch.empty_like(q)
|
||||
|
||||
module.run_kernel(
|
||||
q, k, v, output, qo, kv, batch_size, max_len, 32, 4, 128, 128, 1
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
reference = torch.empty_like(output)
|
||||
import flashinfer
|
||||
workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda")
|
||||
wrapper = flashinfer.BatchPrefillWithRaggedKVCacheWrapper(
|
||||
workspace, kv_layout="NHD", backend="auto"
|
||||
)
|
||||
wrapper.plan(
|
||||
qo, kv, 32, 4, 128, 128, causal=True,
|
||||
q_data_type=torch.bfloat16, kv_data_type=torch.bfloat16,
|
||||
)
|
||||
wrapper.run(q, k, v, out=reference)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
ref_float = reference.float()
|
||||
diff = (output.float() - ref_float).abs()
|
||||
tolerance = ATOL + RTOL * ref_float.abs()
|
||||
match = (diff <= tolerance).float().mean().item()
|
||||
max_abs = diff.max().item()
|
||||
worst_ratio = (diff / tolerance).max().item()
|
||||
required_match = 1.0 if case_id in (14, 15) else 0.99
|
||||
has_nan = torch.isnan(output).any().item()
|
||||
has_inf = torch.isinf(output).any().item()
|
||||
passed = (
|
||||
not has_nan and not has_inf
|
||||
and match >= required_match
|
||||
and worst_ratio <= 8.0
|
||||
)
|
||||
|
||||
print(
|
||||
f"{case_id},{batch_size},{total_q},{total_kv},{max_len},"
|
||||
f"{match:.8f},{max_abs:.8f},{worst_ratio:.8f},"
|
||||
f"{has_nan},{has_inf},{passed}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
if not passed:
|
||||
print(f" FAILED case {case_id}: required_match={required_match}, got {match}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,725 @@
|
|||
import tilelang
|
||||
import tilelang.language as T
|
||||
from tilelang import jit
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 通用配置
|
||||
# ---------------------------------------------------------------------------
|
||||
# 小序列回退 kernel 每个 CTA 处理 BLOCK_M 个 query position。
|
||||
BLOCK_M = 64
|
||||
# seq_len <= 128 时使用的 KV tile 宽度。
|
||||
BLOCK_N = 64
|
||||
# 较长回退形状使用更窄的 KV tile，降低 fragment 和 shared memory 压力。
|
||||
GENERAL_BLOCK_N = 32
|
||||
# 小序列回退 kernel 的线程数；主 packed kernel 会在构建时选择 256 或 512。
|
||||
NUM_THREADS = 128
|
||||
# 当前只使用单阶段循环。K 和 V 在同一块 shared memory 中分阶段复用，
|
||||
# 不能在没有重新设计双缓冲的情况下直接把这里改成 2。
|
||||
NUM_STAGES = 1
|
||||
# softmax 数学形式使用 exp，但设备上的 exp2 指令通常更高效，因此将
|
||||
# exp(x) 转换为 exp2(x * log2(e))。
|
||||
LOG2_E = 1.44269504
|
||||
# TileLang 会针对完整参数组合生成专用 kernel。缓存用于避免同一进程中相同
|
||||
# shape 的重复调用再次触发 JIT 编译；编译时间不应进入 kernel 性能测量。
|
||||
_kernel_cache = {}
|
||||
|
||||
|
||||
@jit(
|
||||
execution_backend="cython",
|
||||
pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True},
|
||||
)
|
||||
def build_kernel(
|
||||
total_q,
|
||||
total_kv,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""构建小序列回退 kernel：每个 CTA 独立处理一个 Q head。"""
|
||||
|
||||
# GQA 中多个 Q head 共享一个 KV head。本题 32 个 Q head、4 个 KV head，
|
||||
# 因此 group_size=8，Q head h 对应的 KV head 为 h // 8。
|
||||
group_size = num_qo_heads // num_kv_heads
|
||||
# 标准 attention scale 为 1/sqrt(head_dim_qk)，再乘 log2(e) 供 exp2 使用。
|
||||
softmax_scale = (1.0 / head_dim_qk) ** 0.5 * LOG2_E
|
||||
# 极短序列使用 N=64，减少 KV 循环次数；较长回退形状使用 N=32，减少
|
||||
# scores/probs fragment 和 shared memory 的瞬时占用。
|
||||
block_n = BLOCK_N if seq_len <= 128 else GENERAL_BLOCK_N
|
||||
# seq_len 是每个 ragged 请求长度的上界。若 total_q=batch_size*seq_len，
|
||||
# 所有请求都只能恰好等于 seq_len，此时网格中不存在无效 Q tile，可以让
|
||||
# JIT 在编译期删除 valid_q_tile 对循环范围的保护逻辑。
|
||||
guard_invalid_tiles = total_q != batch_size * seq_len
|
||||
|
||||
@T.prim_func
|
||||
def kernel(
|
||||
q: T.Tensor((total_q, num_qo_heads, head_dim_qk), T.bfloat16),
|
||||
k: T.Tensor((total_kv, num_kv_heads, head_dim_qk), T.bfloat16),
|
||||
v: T.Tensor((total_kv, num_kv_heads, head_dim_vo), T.bfloat16),
|
||||
output: T.Tensor((total_q, num_qo_heads, head_dim_vo), T.bfloat16),
|
||||
qo_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
kv_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
):
|
||||
# 网格维度：
|
||||
# x = 一个请求内的 Q tile 编号；
|
||||
# y = Q head 编号；
|
||||
# z = batch/request 编号。
|
||||
# 因此一个 CTA 负责 (batch_idx, qo_head, q_tile) 的完整 attention。
|
||||
with T.Kernel(
|
||||
T.ceildiv(seq_len, BLOCK_M),
|
||||
num_qo_heads,
|
||||
batch_size,
|
||||
threads=NUM_THREADS,
|
||||
) as (q_tile, qo_head, batch_idx):
|
||||
# shared memory 保存本 CTA 重复使用的 Q/K/V tile。Q 在整个 KV
|
||||
# 循环中保持不变，K 和 V 则随 kv_tile 更新。
|
||||
q_shared = T.alloc_shared((BLOCK_M, head_dim_qk), T.bfloat16)
|
||||
k_shared = T.alloc_shared((block_n, head_dim_qk), T.bfloat16)
|
||||
v_shared = T.alloc_shared((block_n, head_dim_vo), T.bfloat16)
|
||||
|
||||
# fragment 通常映射到线程私有寄存器/矩阵累加器：
|
||||
# scores : 当前 Q tile 与当前 K tile 的 FP32 logits；
|
||||
# probs : 将 softmax 权重转成 BF16，作为 PV GEMM 的输入；
|
||||
# output_accum : 跨所有 KV tile 累积的 FP32 输出分子。
|
||||
scores = T.alloc_fragment((BLOCK_M, block_n), T.float32)
|
||||
probs = T.alloc_fragment((BLOCK_M, block_n), T.bfloat16)
|
||||
output_accum = T.alloc_fragment((BLOCK_M, head_dim_vo), T.float32)
|
||||
# 在线 Softmax 每一行只维护少量状态：
|
||||
# row_max : 截止当前 KV tile 的全局最大 logit；
|
||||
# row_max_prev : 临时 scratch，先存当前 tile 最大值，后存指数和；
|
||||
# row_scale : 最大值变化后，历史累加结果需要乘的缩放因子；
|
||||
# row_denom : 截止当前 KV tile 的 softmax 分母。
|
||||
row_max = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_max_prev = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_scale = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_denom = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
|
||||
# indptr 给出当前 ragged 请求在扁平 Q/K/V 张量中的区间。
|
||||
q_start = qo_indptr[batch_idx]
|
||||
q_end = qo_indptr[batch_idx + 1]
|
||||
kv_start = kv_indptr[batch_idx]
|
||||
kv_end = kv_indptr[batch_idx + 1]
|
||||
q_len = q_end - q_start
|
||||
kv_len = kv_end - kv_start
|
||||
# 根据 GQA 分组找到该 Q head 共享的 KV head。
|
||||
kv_head = qo_head // group_size
|
||||
# bottom-right causal 对齐的偏移量。可见条件为：
|
||||
# kv_pos < q_pos + 1 + (kv_len - q_len)。
|
||||
causal_offset = kv_len - q_len
|
||||
# 网格按全局 seq_len 上界启动，ragged 请求可能没有对应的 q_tile。
|
||||
valid_q_tile = (
|
||||
q_tile * BLOCK_M < q_len if guard_invalid_tiles else True
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# Q tile 在整个 KV 循环中都会复用，所以只在循环前加载一次。
|
||||
T.copy(
|
||||
q[
|
||||
q_start + q_tile * BLOCK_M : q_start + (q_tile + 1) * BLOCK_M,
|
||||
qo_head,
|
||||
:,
|
||||
],
|
||||
q_shared,
|
||||
)
|
||||
# 在线 Softmax 初始状态：输出分子和分母为 0，最大值为 -inf。
|
||||
T.fill(output_accum, 0)
|
||||
T.fill(row_denom, 0)
|
||||
T.fill(row_max, -T.infinity(T.float32))
|
||||
|
||||
# q_tile 中最后一行 query 最多能看到的 KV 长度。提前缩短 KV 循环，
|
||||
# 避免对因果边界右侧完全不可见的 KV tile 执行 GEMM。
|
||||
max_visible = T.min(
|
||||
T.max(0, causal_offset + (q_tile + 1) * BLOCK_M),
|
||||
kv_len,
|
||||
)
|
||||
visible_tile_count = (
|
||||
T.ceildiv(max_visible, block_n)
|
||||
if causal
|
||||
else T.ceildiv(kv_len, block_n)
|
||||
)
|
||||
# 对无效 q_tile 将循环次数设为 0，从而跳过 K/V copy、两个 GEMM 和
|
||||
# Softmax，而不仅仅是在最后禁止写回。
|
||||
loop_range = (
|
||||
T.if_then_else(valid_q_tile, visible_tile_count, 0)
|
||||
if guard_invalid_tiles
|
||||
else visible_tile_count
|
||||
)
|
||||
|
||||
# 逐块扫描当前 Q tile 可见的 KV 区域。NUM_STAGES=1 表示这里没有
|
||||
# 跨 kv_tile 的 K/V 双缓冲，所有在线 Softmax 状态都存在循环依赖。
|
||||
for kv_tile in T.Pipelined(loop_range, num_stages=NUM_STAGES):
|
||||
# 加载当前 K tile。回退路径主要处理短序列，边界安全由 TileLang
|
||||
# 对 copy 的合法范围处理以及后续显式 mask 共同保证。
|
||||
T.copy(
|
||||
k[
|
||||
kv_start + kv_tile * block_n : kv_start + (kv_tile + 1) * block_n,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
k_shared,
|
||||
)
|
||||
|
||||
# 先把 scores 初始化为 mask：合法元素为 0，非法元素为大负数。
|
||||
# 后续 QK GEMM 默认累加到 scores，因此最终得到 QK 或被 mask 的
|
||||
# 大负数，而不是让 GEMM 覆盖掉这里的因果/尾部 mask。
|
||||
for i, j in T.Parallel(BLOCK_M, block_n):
|
||||
q_pos = q_tile * BLOCK_M + i
|
||||
kv_pos = kv_tile * block_n + j
|
||||
scores[i, j] = T.if_then_else(
|
||||
(q_pos >= q_len)
|
||||
or (kv_pos >= kv_len)
|
||||
or (causal and kv_pos >= q_pos + 1 + causal_offset),
|
||||
-1.0e9,
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 计算 Q @ K^T。默认 clear_accum=False，会保留上面写入 scores
|
||||
# 的初始 mask 并将矩阵乘结果累加进去。
|
||||
T.gemm(
|
||||
q_shared,
|
||||
k_shared,
|
||||
scores,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
# 第一步：求当前 KV tile 每一行的最大值，暂存在 row_max_prev。
|
||||
# row_max 仍然保存此前所有 KV tile 的 running max。
|
||||
T.reduce_max(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
# 若当前 tile 提高了最大值，历史分母和历史输出分子都必须乘：
|
||||
# exp(old_max-new_max)。这里已换算为 exp2 域。
|
||||
row_scale[i] = T.exp2(
|
||||
row_max[i] * softmax_scale
|
||||
- T.max(row_max[i], row_max_prev[i])
|
||||
* softmax_scale
|
||||
)
|
||||
# 将 running max 更新为包含当前 tile 的新最大值。
|
||||
row_max[i] = T.max(row_max[i], row_max_prev[i])
|
||||
|
||||
# 第二步：计算当前 tile 相对于新 running max 的未归一化指数值。
|
||||
# 此处不立即除以分母，避免每个 KV tile 都执行完整归一化。
|
||||
for i, j in T.Parallel(BLOCK_M, block_n):
|
||||
scores[i, j] = T.exp2(
|
||||
scores[i, j] * softmax_scale
|
||||
- row_max[i] * softmax_scale
|
||||
)
|
||||
# row_max_prev 的 tile-max 已完成使命，现在复用同一个 fragment
|
||||
# 保存当前 tile 每一行的指数和，从而避免单独分配 row_sum。
|
||||
T.reduce_sum(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
# 更新在线 Softmax 分母：先把历史分母调整到新最大值尺度，
|
||||
# 再加上当前 tile 的指数和。
|
||||
row_denom[i] = (
|
||||
row_denom[i] * row_scale[i] + row_max_prev[i]
|
||||
)
|
||||
# PV GEMM 使用 BF16 输入，因此把 FP32 指数权重转换到 probs。
|
||||
T.copy(scores, probs)
|
||||
|
||||
# 输出分子和分母必须处于相同的最大值尺度。先缩放历史输出分子，
|
||||
# 再通过下面的 PV GEMM 加入当前 tile 的贡献。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 当前 K 已经不再使用，加载对应的 V tile。
|
||||
T.copy(
|
||||
v[
|
||||
kv_start + kv_tile * block_n : kv_start + (kv_tile + 1) * block_n,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
v_shared,
|
||||
)
|
||||
# output_accum += probs @ V，保持 FP32 累加精度。
|
||||
T.gemm(
|
||||
probs,
|
||||
v_shared,
|
||||
output_accum,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# KV 循环结束后才做最终归一化。每行只计算一次 1/denom，避免
|
||||
# head_dim_vo 个输出元素分别执行相同的除法。
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
q_pos = q_tile * BLOCK_M + i
|
||||
row_scale[i] = T.if_then_else(
|
||||
(q_pos < q_len)
|
||||
and (
|
||||
(not causal)
|
||||
or q_pos + causal_offset >= 0
|
||||
),
|
||||
1.0 / row_denom[i],
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 用乘法完成整行归一化。完全被 mask 的行乘 0。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 只写回请求实际存在的 query 行，抑制最后一个 Q tile 的尾部。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
if q_tile * BLOCK_M + i < q_len:
|
||||
output[
|
||||
q_start + q_tile * BLOCK_M + i,
|
||||
qo_head,
|
||||
d,
|
||||
] = output_accum[i, d]
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
@jit(
|
||||
execution_backend="cython",
|
||||
pass_configs={
|
||||
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
|
||||
tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
|
||||
},
|
||||
compile_flags=["-O3", "-DENABLE_BF16"],
|
||||
)
|
||||
def build_packed_kernel_pv_fullcol_v14(
|
||||
total_q,
|
||||
total_kv,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""构建主 GQA kernel，并在长路径按输出列划分 PV waves。"""
|
||||
|
||||
# equal-length 且 seq_len<=1024 时使用 M64/N32/256 threads。较小的资源
|
||||
# 占用可提高 CTA 驻留数；长序列和非对称形状保持 M128/N64/512 threads，
|
||||
# 以减少 Q tile 数量、KV 循环次数和在线 Softmax 更新次数。
|
||||
use_resident2 = total_q == total_kv and seq_len <= 1024
|
||||
packed_block_m = 64 if use_resident2 else 128
|
||||
packed_block_n = 32 if use_resident2 else 64
|
||||
packed_threads = 256 if use_resident2 else 512
|
||||
# opt_012 的 case-4 生成代码中，FullRow PV 会把 V 从 shared memory
|
||||
# 逐个 BF16 标量搬到 B fragment；profiler 同时显示 shared efficiency
|
||||
# 只有 74.04%，VLS pipeline stall 占比升至 10.72%。长路径改用 FullCol，
|
||||
# 让 waves 沿 128 维输出列划分，目标是生成更连续的 V shared 读取。
|
||||
# resident2 短路径保留已验证的 FullRow 映射，避免扩大实验影响面。
|
||||
pv_policy = (
|
||||
T.GemmWarpPolicy.FullRow
|
||||
if use_resident2
|
||||
else T.GemmWarpPolicy.FullCol
|
||||
)
|
||||
# 本题 group_size=8。一个 KV head 对应连续的 8 个 Q head。
|
||||
group_size = num_qo_heads // num_kv_heads
|
||||
# packed M 维中的每一行不是单纯的 query position，而是：
|
||||
# packed_row = q_pos * group_size + group_head。
|
||||
# 因此同一个 CTA 可以处理一个 KV head 对应的多个 Q head，并让这些 Q
|
||||
# head 在 CTA 内共享 K/V 的 global->shared 加载。
|
||||
packed_tile_count = T.ceildiv(seq_len * group_size, packed_block_m)
|
||||
# dense_equal 是编译期常量。由于每个 segment 长度都不超过 seq_len，而
|
||||
# 总长度恰好为 batch_size*seq_len，可以严格推出每个 segment 都等长。
|
||||
dense_equal = (
|
||||
total_q == batch_size * seq_len
|
||||
and total_kv == batch_size * seq_len
|
||||
)
|
||||
# 对 dense 长序列启用反向逻辑 tile 映射。因果 attention 中越靠后的 Q
|
||||
# tile 能看到越多 KV，计算量越大；优先提交重 tile 可以缩短最后少数 CTA
|
||||
# 造成的调度长尾。resident2 和 ragged 路径仍保持正向映射。
|
||||
reverse_dense_tiles = dense_equal and not use_resident2
|
||||
# attention scale 转为 exp2 所需的 base-2 缩放系数。
|
||||
softmax_scale = (1.0 / head_dim_qk) ** 0.5 * LOG2_E
|
||||
|
||||
@T.prim_func
|
||||
def packed_kernel_pv_fullcol_v14(
|
||||
q: T.Tensor((total_q, num_qo_heads, head_dim_qk), T.bfloat16),
|
||||
k: T.Tensor((total_kv, num_kv_heads, head_dim_qk), T.bfloat16),
|
||||
v: T.Tensor((total_kv, num_kv_heads, head_dim_vo), T.bfloat16),
|
||||
output: T.Tensor((total_q, num_qo_heads, head_dim_vo), T.bfloat16),
|
||||
qo_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
kv_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
):
|
||||
# 网格维度：
|
||||
# x = packed Q tile；
|
||||
# y = KV head；
|
||||
# z = batch/request。
|
||||
# 一个 CTA 因而处理一个请求、一个 KV head 和一段 packed Q rows。
|
||||
with T.Kernel(
|
||||
packed_tile_count,
|
||||
num_kv_heads,
|
||||
batch_size,
|
||||
threads=packed_threads,
|
||||
) as (packed_tile, kv_head, batch_idx):
|
||||
# Q tile 在整个 KV 循环中常驻 shared memory。长路径 M=128 时，
|
||||
# q_shared 大小为 128*128*2 = 32 KiB。
|
||||
q_shared = T.alloc_shared(
|
||||
(packed_block_m, head_dim_qk), T.bfloat16
|
||||
)
|
||||
# 同一个 kv_shared 分时保存 K 和 V：QK GEMM 完成后当前 K 生命周期
|
||||
# 已结束，随后 V 覆盖这块空间。这样长路径只需 16 KiB KV shared，
|
||||
# 而不是分别为 K、V 分配两份内存。
|
||||
kv_shared = T.alloc_shared(
|
||||
(packed_block_n, head_dim_qk), T.bfloat16
|
||||
)
|
||||
|
||||
# scores：FP32 QK 累加器和当前 tile 的指数权重。
|
||||
scores = T.alloc_fragment(
|
||||
(packed_block_m, packed_block_n), T.float32
|
||||
)
|
||||
# probs：scores 转成 BF16 后供 PV 矩阵乘使用。
|
||||
probs = T.alloc_fragment(
|
||||
(packed_block_m, packed_block_n), T.bfloat16
|
||||
)
|
||||
# output_accum：跨全部 KV tile 保存 FP32 输出分子。
|
||||
output_accum = T.alloc_fragment(
|
||||
(packed_block_m, head_dim_vo), T.float32
|
||||
)
|
||||
# 每个 packed row 的在线 Softmax 状态。row_max_prev 是 scratch，
|
||||
# row_scale 在循环中保存历史尺度修正，循环结束后复用为 1/denom。
|
||||
row_max = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_max_prev = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_scale = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_denom = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
|
||||
# dense equal-length 特化：直接用 batch_idx*seq_len 定位请求，并将
|
||||
# q_len/kv_len 变成编译期常量；ragged 路径仍从 indptr 读取真实边界。
|
||||
# 这可删除 dense 热路径上的 indptr load 和部分动态边界判断。
|
||||
q_start = (
|
||||
batch_idx * seq_len
|
||||
if dense_equal
|
||||
else qo_indptr[batch_idx]
|
||||
)
|
||||
kv_start = (
|
||||
batch_idx * seq_len
|
||||
if dense_equal
|
||||
else kv_indptr[batch_idx]
|
||||
)
|
||||
q_len = (
|
||||
seq_len
|
||||
if dense_equal
|
||||
else qo_indptr[batch_idx + 1] - q_start
|
||||
)
|
||||
kv_len = (
|
||||
seq_len
|
||||
if dense_equal
|
||||
else kv_indptr[batch_idx + 1] - kv_start
|
||||
)
|
||||
# MetaX PipelinePlanning 在 dense 长度完全常量化后，会把同一
|
||||
# kv_shared 中先写 K、后写 V 的合法生命周期复用误判成 stage
|
||||
# 重叠写。dense 路径仅为 copy 边界保留一次运行时 indptr 读取，
|
||||
# 数学边界、网格和有效 tile 判断仍然使用常量 kv_len。
|
||||
copy_kv_len = (
|
||||
kv_indptr[batch_idx + 1] - kv_start
|
||||
if dense_equal
|
||||
else kv_len
|
||||
)
|
||||
# packed_q_len 是当前请求包含的逻辑 packed rows 数量。
|
||||
packed_q_len = q_len * group_size
|
||||
# physical packed_tile 来自 blockIdx.x；logical_packed_tile 决定实际
|
||||
# 处理哪段 Q。反转只改变 CTA 提交顺序，不改变数学结果或输出位置。
|
||||
logical_packed_tile = (
|
||||
packed_tile_count - 1 - packed_tile
|
||||
if reverse_dense_tiles
|
||||
else packed_tile
|
||||
)
|
||||
packed_tile_start = logical_packed_tile * packed_block_m
|
||||
# dense 路径的网格恰好覆盖全部 packed rows，因此所有 CTA 有效；
|
||||
# ragged 请求可能短于 seq_len 上界，需要在运行时过滤无效 CTA。
|
||||
valid_q_tile = (
|
||||
True
|
||||
if dense_equal
|
||||
else packed_tile_start < packed_q_len
|
||||
)
|
||||
# bottom-right causal mask 使用 kv_len-q_len 修正 Q/KV 长度差。
|
||||
causal_offset = kv_len - q_len
|
||||
|
||||
if valid_q_tile:
|
||||
# 将 packed row 反解为 query position 和 GQA 组内 Q head：
|
||||
# q_pos = packed_row // group_size
|
||||
# group_head= packed_row % group_size
|
||||
# 再与当前 kv_head 组合出原始 qo_head。最后一个 tile 的补齐行写0。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_qk):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
group_head = packed_row % group_size
|
||||
if packed_row < packed_q_len:
|
||||
q_shared[i, d] = q[
|
||||
q_start + q_pos,
|
||||
kv_head * group_size + group_head,
|
||||
d,
|
||||
]
|
||||
else:
|
||||
q_shared[i, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
# 初始化在线 Softmax：输出分子=0、分母=0、running max=-inf。
|
||||
T.clear(output_accum)
|
||||
T.clear(row_denom)
|
||||
T.fill(row_max, -T.infinity(T.float32))
|
||||
|
||||
# 当前 packed tile 末尾对应的 query position 上界。一个 packed tile
|
||||
# 覆盖 packed_block_m/group_size 个不同 query position。
|
||||
q_upper_bound = T.min(
|
||||
q_len,
|
||||
T.ceildiv(
|
||||
(logical_packed_tile + 1) * packed_block_m,
|
||||
group_size,
|
||||
),
|
||||
)
|
||||
# 使用 tile 内最后一个 query 的因果边界求最大可见 KV 长度，以便
|
||||
# 整块跳过其右侧不可能被任何行看到的 KV tiles。
|
||||
max_visible = T.min(
|
||||
T.max(0, causal_offset + q_upper_bound),
|
||||
kv_len,
|
||||
)
|
||||
visible_tile_count = (
|
||||
T.ceildiv(max_visible, packed_block_n)
|
||||
if causal
|
||||
else T.ceildiv(kv_len, packed_block_n)
|
||||
)
|
||||
# 无效 ragged Q tile 的 loop_range=0，可跳过全部主要计算。
|
||||
loop_range = T.if_then_else(
|
||||
valid_q_tile, visible_tile_count, 0
|
||||
)
|
||||
# 当前 packed tile 第一行对应的 query position。
|
||||
first_q_pos = packed_tile_start // group_size
|
||||
# 如果某个 KV tile 连当前 packed tile 的第一行都完全可见，那么它
|
||||
# 对后续所有行也完全可见。这部分 tile 无需逐元素生成 causal mask，
|
||||
# 直接把 scores 清零作为 GEMM 初始累加器即可。
|
||||
fully_visible_tiles = T.min(
|
||||
loop_range,
|
||||
T.max(
|
||||
0,
|
||||
(first_q_pos + causal_offset + 1) // packed_block_n,
|
||||
),
|
||||
)
|
||||
|
||||
# 顺序扫描当前 packed Q tile 可见的 KV tiles。在线 Softmax 的
|
||||
# running max/denom/output_accum 在相邻迭代间存在严格依赖。
|
||||
for kv_tile in T.Pipelined(loop_range, num_stages=NUM_STAGES):
|
||||
# 当前 KV tile 在请求内部的半开区间 [tile_start, tile_end)。
|
||||
tile_start = kv_tile * packed_block_n
|
||||
tile_end = tile_start + packed_block_n
|
||||
|
||||
if tile_end <= copy_kv_len:
|
||||
# 完整 K tile 走向量化 copy。disable_tma=True 使用当前
|
||||
# MetaX 后端已验证可工作的普通 shared-memory copy 路径。
|
||||
T.copy(
|
||||
k[
|
||||
kv_start + tile_start : kv_start + tile_end,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
kv_shared,
|
||||
disable_tma=True,
|
||||
)
|
||||
else:
|
||||
# 最后一个 ragged K tile 可能越过请求边界，必须逐元素判断；
|
||||
# 越界位置填0，不能读到扁平张量中下一个请求的数据。
|
||||
for j, d in T.Parallel(packed_block_n, head_dim_qk):
|
||||
kv_pos = tile_start + j
|
||||
if kv_pos < copy_kv_len:
|
||||
kv_shared[j, d] = k[
|
||||
kv_start + kv_pos, kv_head, d
|
||||
]
|
||||
else:
|
||||
kv_shared[j, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
if causal and kv_tile < fully_visible_tiles:
|
||||
# 整块对所有 query 行可见，只需将 QK 累加器初始化为0。
|
||||
T.clear(scores)
|
||||
else:
|
||||
# causal frontier 或 ragged tail 需要逐元素初始化 mask。
|
||||
# 合法条件同时检查 packed Q tail、KV tail 和 bottom-right
|
||||
# causal 边界；非法元素设为 -inf，使其 softmax 权重为0。
|
||||
for i, j in T.Parallel(
|
||||
packed_block_m, packed_block_n
|
||||
):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
kv_pos = tile_start + j
|
||||
scores[i, j] = T.if_then_else(
|
||||
(packed_row < packed_q_len)
|
||||
and (kv_pos < kv_len)
|
||||
and (
|
||||
(not causal)
|
||||
or kv_pos < q_pos + 1 + causal_offset
|
||||
),
|
||||
0.0,
|
||||
-T.infinity(T.float32),
|
||||
)
|
||||
|
||||
# scores = mask + Q @ K^T。clear_accum 默认为 False，因此前面
|
||||
# 写入的 0/-inf mask 会作为矩阵乘累加器被保留下来。
|
||||
T.gemm(
|
||||
q_shared,
|
||||
kv_shared,
|
||||
scores,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
# 求当前 tile 的逐行最大值。row_max_prev 是 scratch，row_max
|
||||
# 始终保存此前全部 KV tiles 的 running max。
|
||||
T.reduce_max(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(packed_block_m):
|
||||
# 最大值改变时，用 exp(old_max-new_max) 把历史分母和历史
|
||||
# 输出分子转换到新的数值尺度，保证在线 Softmax 数值稳定。
|
||||
row_scale[i] = T.exp2(
|
||||
row_max[i] * softmax_scale
|
||||
- T.max(row_max[i], row_max_prev[i])
|
||||
* softmax_scale
|
||||
)
|
||||
# 合并历史最大值与当前 tile 最大值。
|
||||
row_max[i] = T.max(row_max[i], row_max_prev[i])
|
||||
|
||||
# 将 logits 转成相对于新 running max 的指数值。采用
|
||||
# FlashAttention 风格在线 Softmax，不保存完整 attention 矩阵。
|
||||
for i, j in T.Parallel(packed_block_m, packed_block_n):
|
||||
scores[i, j] = T.exp2(
|
||||
scores[i, j] * softmax_scale
|
||||
- row_max[i] * softmax_scale
|
||||
)
|
||||
# tile max 已无后续用途，复用 row_max_prev 保存当前指数和。
|
||||
T.reduce_sum(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(packed_block_m):
|
||||
# 分母更新：历史分母先乘 row_scale，再加当前 tile 指数和。
|
||||
row_denom[i] = (
|
||||
row_denom[i] * row_scale[i] + row_max_prev[i]
|
||||
)
|
||||
# PV GEMM 输入为 BF16，将 FP32 指数权重转换到 probs fragment。
|
||||
T.copy(scores, probs)
|
||||
|
||||
# 历史输出分子与分母使用相同的 row_scale 重新定标。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
if tile_end <= copy_kv_len:
|
||||
# QK 已结束，当前 K 不再存活；把对应 V tile 覆盖写入同一个
|
||||
# kv_shared，以减少每 CTA 的动态 shared memory 占用。
|
||||
T.copy(
|
||||
v[
|
||||
kv_start + tile_start : kv_start + tile_end,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
kv_shared,
|
||||
disable_tma=True,
|
||||
)
|
||||
else:
|
||||
# ragged V tail 与 K tail 一样执行显式边界判断和补0。
|
||||
for j, d in T.Parallel(packed_block_n, head_dim_vo):
|
||||
kv_pos = tile_start + j
|
||||
if kv_pos < copy_kv_len:
|
||||
kv_shared[j, d] = v[
|
||||
kv_start + kv_pos, kv_head, d
|
||||
]
|
||||
else:
|
||||
kv_shared[j, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
# 累积当前 tile 的输出贡献：output_accum += probs @ V。
|
||||
T.gemm(
|
||||
probs,
|
||||
kv_shared,
|
||||
output_accum,
|
||||
policy=pv_policy,
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# 所有 KV tiles 处理完后再归一化。row_scale 已结束循环内使命，
|
||||
# 现在复用为每行的 1/softmax_denominator。每行只除一次，后续
|
||||
# 128 个输出维度都通过乘法归一化。补齐行和完全 mask 行设为0。
|
||||
for i in T.Parallel(packed_block_m):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
row_scale[i] = T.if_then_else(
|
||||
(packed_row < packed_q_len)
|
||||
and (
|
||||
(not causal)
|
||||
or q_pos + causal_offset >= 0
|
||||
),
|
||||
1.0 / row_denom[i],
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 将 packed row 重新映射到原始 output[q_pos, qo_head, d]。
|
||||
# packed tail 只参与内部补齐，不允许写回输出张量。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_vo):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
group_head = packed_row % group_size
|
||||
if packed_row < packed_q_len:
|
||||
output[
|
||||
q_start + q_pos,
|
||||
kv_head * group_size + group_head,
|
||||
d,
|
||||
] = output_accum[i, d] * row_scale[i]
|
||||
|
||||
return packed_kernel_pv_fullcol_v14
|
||||
|
||||
|
||||
def run_kernel(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""OJ 调用入口：按完整规格选择、缓存并启动 TileLang kernel。"""
|
||||
|
||||
# cache key 必须包含所有会改变生成代码、网格或张量 shape 的参数：
|
||||
# q/k 的总长度决定静态张量形状；seq_len 决定网格上界、tile 配置和分支；
|
||||
# head 数、head_dim、causal 则直接影响 GQA 映射和 attention 数学逻辑。
|
||||
key = (
|
||||
"packed-gqa-pv-fullcol-v14",
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
if key not in _kernel_cache:
|
||||
# 首次遇到该参数组合时才执行 JIT：
|
||||
# seq_len<=128：小序列回退路径，每 CTA 处理一个 Q head，启动和资源
|
||||
# 开销更小，保留单 token、非2次幂尾部的现有优势；
|
||||
# seq_len>128 ：主 packed GQA 路径，在一个 CTA 内处理同一 KV head
|
||||
# 对应的多个 Q heads，提高 K/V 数据复用。
|
||||
if seq_len <= 128:
|
||||
_kernel_cache[key] = build_kernel(
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
else:
|
||||
_kernel_cache[key] = build_packed_kernel_pv_fullcol_v14(
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
# JIT 结果是可直接接收 Torch tensors 的 Cython backend callable。
|
||||
# output 由调用方提前分配，本函数只启动 kernel，不在计时路径创建临时张量。
|
||||
_kernel_cache[key](q, k, v, output, qo_indptr, kv_indptr)
|
||||
|
|
@ -0,0 +1,719 @@
|
|||
import tilelang
|
||||
import tilelang.language as T
|
||||
from tilelang import jit
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 通用配置
|
||||
# ---------------------------------------------------------------------------
|
||||
# 小序列回退 kernel 每个 CTA 处理 BLOCK_M 个 query position。
|
||||
BLOCK_M = 64
|
||||
# seq_len <= 128 时使用的 KV tile 宽度。
|
||||
BLOCK_N = 64
|
||||
# 较长回退形状使用更窄的 KV tile，降低 fragment 和 shared memory 压力。
|
||||
GENERAL_BLOCK_N = 32
|
||||
# 小序列回退 kernel 的线程数；主 packed kernel 会在构建时选择 256 或 512。
|
||||
NUM_THREADS = 128
|
||||
# 当前只使用单阶段循环。K 和 V 在同一块 shared memory 中分阶段复用，
|
||||
# 不能在没有重新设计双缓冲的情况下直接把这里改成 2。
|
||||
NUM_STAGES = 1
|
||||
# softmax 数学形式使用 exp，但设备上的 exp2 指令通常更高效，因此将
|
||||
# exp(x) 转换为 exp2(x * log2(e))。
|
||||
LOG2_E = 1.44269504
|
||||
# TileLang 会针对完整参数组合生成专用 kernel。缓存用于避免同一进程中相同
|
||||
# shape 的重复调用再次触发 JIT 编译；编译时间不应进入 kernel 性能测量。
|
||||
_kernel_cache = {}
|
||||
|
||||
|
||||
@jit(
|
||||
execution_backend="cython",
|
||||
pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True},
|
||||
)
|
||||
def build_kernel(
|
||||
total_q,
|
||||
total_kv,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""构建小序列回退 kernel：每个 CTA 独立处理一个 Q head。"""
|
||||
|
||||
# GQA 中多个 Q head 共享一个 KV head。本题 32 个 Q head、4 个 KV head，
|
||||
# 因此 group_size=8，Q head h 对应的 KV head 为 h // 8。
|
||||
group_size = num_qo_heads // num_kv_heads
|
||||
# 标准 attention scale 为 1/sqrt(head_dim_qk)，再乘 log2(e) 供 exp2 使用。
|
||||
softmax_scale = (1.0 / head_dim_qk) ** 0.5 * LOG2_E
|
||||
# 极短序列使用 N=64，减少 KV 循环次数；较长回退形状使用 N=32，减少
|
||||
# scores/probs fragment 和 shared memory 的瞬时占用。
|
||||
block_n = BLOCK_N if seq_len <= 128 else GENERAL_BLOCK_N
|
||||
# seq_len 是每个 ragged 请求长度的上界。若 total_q=batch_size*seq_len，
|
||||
# 所有请求都只能恰好等于 seq_len，此时网格中不存在无效 Q tile，可以让
|
||||
# JIT 在编译期删除 valid_q_tile 对循环范围的保护逻辑。
|
||||
guard_invalid_tiles = total_q != batch_size * seq_len
|
||||
|
||||
@T.prim_func
|
||||
def kernel(
|
||||
q: T.Tensor((total_q, num_qo_heads, head_dim_qk), T.bfloat16),
|
||||
k: T.Tensor((total_kv, num_kv_heads, head_dim_qk), T.bfloat16),
|
||||
v: T.Tensor((total_kv, num_kv_heads, head_dim_vo), T.bfloat16),
|
||||
output: T.Tensor((total_q, num_qo_heads, head_dim_vo), T.bfloat16),
|
||||
qo_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
kv_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
):
|
||||
# 网格维度：
|
||||
# x = 一个请求内的 Q tile 编号；
|
||||
# y = Q head 编号；
|
||||
# z = batch/request 编号。
|
||||
# 因此一个 CTA 负责 (batch_idx, qo_head, q_tile) 的完整 attention。
|
||||
with T.Kernel(
|
||||
T.ceildiv(seq_len, BLOCK_M),
|
||||
num_qo_heads,
|
||||
batch_size,
|
||||
threads=NUM_THREADS,
|
||||
) as (q_tile, qo_head, batch_idx):
|
||||
# shared memory 保存本 CTA 重复使用的 Q/K/V tile。Q 在整个 KV
|
||||
# 循环中保持不变，K 和 V 则随 kv_tile 更新。
|
||||
q_shared = T.alloc_shared((BLOCK_M, head_dim_qk), T.bfloat16)
|
||||
k_shared = T.alloc_shared((block_n, head_dim_qk), T.bfloat16)
|
||||
v_shared = T.alloc_shared((block_n, head_dim_vo), T.bfloat16)
|
||||
|
||||
# fragment 通常映射到线程私有寄存器/矩阵累加器：
|
||||
# scores : 当前 Q tile 与当前 K tile 的 FP32 logits；
|
||||
# probs : 将 softmax 权重转成 BF16，作为 PV GEMM 的输入；
|
||||
# output_accum : 跨所有 KV tile 累积的 FP32 输出分子。
|
||||
scores = T.alloc_fragment((BLOCK_M, block_n), T.float32)
|
||||
probs = T.alloc_fragment((BLOCK_M, block_n), T.bfloat16)
|
||||
output_accum = T.alloc_fragment((BLOCK_M, head_dim_vo), T.float32)
|
||||
# 在线 Softmax 每一行只维护少量状态：
|
||||
# row_max : 截止当前 KV tile 的全局最大 logit；
|
||||
# row_max_prev : 临时 scratch，先存当前 tile 最大值，后存指数和；
|
||||
# row_scale : 最大值变化后，历史累加结果需要乘的缩放因子；
|
||||
# row_denom : 截止当前 KV tile 的 softmax 分母。
|
||||
row_max = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_max_prev = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_scale = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_denom = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
|
||||
# indptr 给出当前 ragged 请求在扁平 Q/K/V 张量中的区间。
|
||||
q_start = qo_indptr[batch_idx]
|
||||
q_end = qo_indptr[batch_idx + 1]
|
||||
kv_start = kv_indptr[batch_idx]
|
||||
kv_end = kv_indptr[batch_idx + 1]
|
||||
q_len = q_end - q_start
|
||||
kv_len = kv_end - kv_start
|
||||
# 根据 GQA 分组找到该 Q head 共享的 KV head。
|
||||
kv_head = qo_head // group_size
|
||||
# bottom-right causal 对齐的偏移量。可见条件为：
|
||||
# kv_pos < q_pos + 1 + (kv_len - q_len)。
|
||||
causal_offset = kv_len - q_len
|
||||
# 网格按全局 seq_len 上界启动，ragged 请求可能没有对应的 q_tile。
|
||||
valid_q_tile = (
|
||||
q_tile * BLOCK_M < q_len if guard_invalid_tiles else True
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# Q tile 在整个 KV 循环中都会复用，所以只在循环前加载一次。
|
||||
T.copy(
|
||||
q[
|
||||
q_start + q_tile * BLOCK_M : q_start + (q_tile + 1) * BLOCK_M,
|
||||
qo_head,
|
||||
:,
|
||||
],
|
||||
q_shared,
|
||||
)
|
||||
# 在线 Softmax 初始状态：输出分子和分母为 0，最大值为 -inf。
|
||||
T.fill(output_accum, 0)
|
||||
T.fill(row_denom, 0)
|
||||
T.fill(row_max, -T.infinity(T.float32))
|
||||
|
||||
# q_tile 中最后一行 query 最多能看到的 KV 长度。提前缩短 KV 循环，
|
||||
# 避免对因果边界右侧完全不可见的 KV tile 执行 GEMM。
|
||||
max_visible = T.min(
|
||||
T.max(0, causal_offset + (q_tile + 1) * BLOCK_M),
|
||||
kv_len,
|
||||
)
|
||||
visible_tile_count = (
|
||||
T.ceildiv(max_visible, block_n)
|
||||
if causal
|
||||
else T.ceildiv(kv_len, block_n)
|
||||
)
|
||||
# 对无效 q_tile 将循环次数设为 0，从而跳过 K/V copy、两个 GEMM 和
|
||||
# Softmax，而不仅仅是在最后禁止写回。
|
||||
loop_range = (
|
||||
T.if_then_else(valid_q_tile, visible_tile_count, 0)
|
||||
if guard_invalid_tiles
|
||||
else visible_tile_count
|
||||
)
|
||||
|
||||
# 逐块扫描当前 Q tile 可见的 KV 区域。NUM_STAGES=1 表示这里没有
|
||||
# 跨 kv_tile 的 K/V 双缓冲，所有在线 Softmax 状态都存在循环依赖。
|
||||
for kv_tile in T.Pipelined(loop_range, num_stages=NUM_STAGES):
|
||||
# 加载当前 K tile。回退路径主要处理短序列，边界安全由 TileLang
|
||||
# 对 copy 的合法范围处理以及后续显式 mask 共同保证。
|
||||
T.copy(
|
||||
k[
|
||||
kv_start + kv_tile * block_n : kv_start + (kv_tile + 1) * block_n,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
k_shared,
|
||||
)
|
||||
|
||||
# 先把 scores 初始化为 mask：合法元素为 0，非法元素为大负数。
|
||||
# 后续 QK GEMM 默认累加到 scores，因此最终得到 QK 或被 mask 的
|
||||
# 大负数，而不是让 GEMM 覆盖掉这里的因果/尾部 mask。
|
||||
for i, j in T.Parallel(BLOCK_M, block_n):
|
||||
q_pos = q_tile * BLOCK_M + i
|
||||
kv_pos = kv_tile * block_n + j
|
||||
scores[i, j] = T.if_then_else(
|
||||
(q_pos >= q_len)
|
||||
or (kv_pos >= kv_len)
|
||||
or (causal and kv_pos >= q_pos + 1 + causal_offset),
|
||||
-1.0e9,
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 计算 Q @ K^T。默认 clear_accum=False，会保留上面写入 scores
|
||||
# 的初始 mask 并将矩阵乘结果累加进去。
|
||||
T.gemm(
|
||||
q_shared,
|
||||
k_shared,
|
||||
scores,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
# 第一步：求当前 KV tile 每一行的最大值，暂存在 row_max_prev。
|
||||
# row_max 仍然保存此前所有 KV tile 的 running max。
|
||||
T.reduce_max(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
# 若当前 tile 提高了最大值，历史分母和历史输出分子都必须乘：
|
||||
# exp(old_max-new_max)。这里已换算为 exp2 域。
|
||||
row_scale[i] = T.exp2(
|
||||
row_max[i] * softmax_scale
|
||||
- T.max(row_max[i], row_max_prev[i])
|
||||
* softmax_scale
|
||||
)
|
||||
# 将 running max 更新为包含当前 tile 的新最大值。
|
||||
row_max[i] = T.max(row_max[i], row_max_prev[i])
|
||||
|
||||
# 第二步：计算当前 tile 相对于新 running max 的未归一化指数值。
|
||||
# 此处不立即除以分母，避免每个 KV tile 都执行完整归一化。
|
||||
for i, j in T.Parallel(BLOCK_M, block_n):
|
||||
scores[i, j] = T.exp2(
|
||||
scores[i, j] * softmax_scale
|
||||
- row_max[i] * softmax_scale
|
||||
)
|
||||
# row_max_prev 的 tile-max 已完成使命，现在复用同一个 fragment
|
||||
# 保存当前 tile 每一行的指数和，从而避免单独分配 row_sum。
|
||||
T.reduce_sum(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
# 更新在线 Softmax 分母：先把历史分母调整到新最大值尺度，
|
||||
# 再加上当前 tile 的指数和。
|
||||
row_denom[i] = (
|
||||
row_denom[i] * row_scale[i] + row_max_prev[i]
|
||||
)
|
||||
# PV GEMM 使用 BF16 输入，因此把 FP32 指数权重转换到 probs。
|
||||
T.copy(scores, probs)
|
||||
|
||||
# 输出分子和分母必须处于相同的最大值尺度。先缩放历史输出分子，
|
||||
# 再通过下面的 PV GEMM 加入当前 tile 的贡献。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 当前 K 已经不再使用，加载对应的 V tile。
|
||||
T.copy(
|
||||
v[
|
||||
kv_start + kv_tile * block_n : kv_start + (kv_tile + 1) * block_n,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
v_shared,
|
||||
)
|
||||
# output_accum += probs @ V，保持 FP32 累加精度。
|
||||
T.gemm(
|
||||
probs,
|
||||
v_shared,
|
||||
output_accum,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# KV 循环结束后才做最终归一化。每行只计算一次 1/denom，避免
|
||||
# head_dim_vo 个输出元素分别执行相同的除法。
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
q_pos = q_tile * BLOCK_M + i
|
||||
row_scale[i] = T.if_then_else(
|
||||
(q_pos < q_len)
|
||||
and (
|
||||
(not causal)
|
||||
or q_pos + causal_offset >= 0
|
||||
),
|
||||
1.0 / row_denom[i],
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 用乘法完成整行归一化。完全被 mask 的行乘 0。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 只写回请求实际存在的 query 行，抑制最后一个 Q tile 的尾部。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
if q_tile * BLOCK_M + i < q_len:
|
||||
output[
|
||||
q_start + q_tile * BLOCK_M + i,
|
||||
qo_head,
|
||||
d,
|
||||
] = output_accum[i, d]
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
@jit(
|
||||
execution_backend="cython",
|
||||
pass_configs={
|
||||
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
|
||||
tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
|
||||
},
|
||||
compile_flags=["-O3", "-DENABLE_BF16"],
|
||||
)
|
||||
def build_packed_kernel_v_prefetch_after_qk_v15(
|
||||
total_q,
|
||||
total_kv,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""构建主 GQA kernel，并在 QK reduction 后提前加载 V。"""
|
||||
|
||||
# equal-length 且 seq_len<=1024 时使用 M64/N32/256 threads。较小的资源
|
||||
# 占用可提高 CTA 驻留数；长序列和非对称形状保持 M128/N64/512 threads，
|
||||
# 以减少 Q tile 数量、KV 循环次数和在线 Softmax 更新次数。
|
||||
use_resident2 = total_q == total_kv and seq_len <= 1024
|
||||
packed_block_m = 64 if use_resident2 else 128
|
||||
packed_block_n = 32 if use_resident2 else 64
|
||||
packed_threads = 256 if use_resident2 else 512
|
||||
# 本题 group_size=8。一个 KV head 对应连续的 8 个 Q head。
|
||||
group_size = num_qo_heads // num_kv_heads
|
||||
# packed M 维中的每一行不是单纯的 query position，而是：
|
||||
# packed_row = q_pos * group_size + group_head。
|
||||
# 因此同一个 CTA 可以处理一个 KV head 对应的多个 Q head，并让这些 Q
|
||||
# head 在 CTA 内共享 K/V 的 global->shared 加载。
|
||||
packed_tile_count = T.ceildiv(seq_len * group_size, packed_block_m)
|
||||
# dense_equal 是编译期常量。由于每个 segment 长度都不超过 seq_len，而
|
||||
# 总长度恰好为 batch_size*seq_len，可以严格推出每个 segment 都等长。
|
||||
dense_equal = (
|
||||
total_q == batch_size * seq_len
|
||||
and total_kv == batch_size * seq_len
|
||||
)
|
||||
# 对 dense 长序列启用反向逻辑 tile 映射。因果 attention 中越靠后的 Q
|
||||
# tile 能看到越多 KV，计算量越大；优先提交重 tile 可以缩短最后少数 CTA
|
||||
# 造成的调度长尾。resident2 和 ragged 路径仍保持正向映射。
|
||||
reverse_dense_tiles = dense_equal and not use_resident2
|
||||
# attention scale 转为 exp2 所需的 base-2 缩放系数。
|
||||
softmax_scale = (1.0 / head_dim_qk) ** 0.5 * LOG2_E
|
||||
|
||||
@T.prim_func
|
||||
def packed_kernel_v_prefetch_after_qk_v15(
|
||||
q: T.Tensor((total_q, num_qo_heads, head_dim_qk), T.bfloat16),
|
||||
k: T.Tensor((total_kv, num_kv_heads, head_dim_qk), T.bfloat16),
|
||||
v: T.Tensor((total_kv, num_kv_heads, head_dim_vo), T.bfloat16),
|
||||
output: T.Tensor((total_q, num_qo_heads, head_dim_vo), T.bfloat16),
|
||||
qo_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
kv_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
):
|
||||
# 网格维度：
|
||||
# x = packed Q tile；
|
||||
# y = KV head；
|
||||
# z = batch/request。
|
||||
# 一个 CTA 因而处理一个请求、一个 KV head 和一段 packed Q rows。
|
||||
with T.Kernel(
|
||||
packed_tile_count,
|
||||
num_kv_heads,
|
||||
batch_size,
|
||||
threads=packed_threads,
|
||||
) as (packed_tile, kv_head, batch_idx):
|
||||
# Q tile 在整个 KV 循环中常驻 shared memory。长路径 M=128 时，
|
||||
# q_shared 大小为 128*128*2 = 32 KiB。
|
||||
q_shared = T.alloc_shared(
|
||||
(packed_block_m, head_dim_qk), T.bfloat16
|
||||
)
|
||||
# 同一个 kv_shared 分时保存 K 和 V：QK GEMM 完成后当前 K 生命周期
|
||||
# 已结束，随后 V 覆盖这块空间。这样长路径只需 16 KiB KV shared，
|
||||
# 而不是分别为 K、V 分配两份内存。
|
||||
kv_shared = T.alloc_shared(
|
||||
(packed_block_n, head_dim_qk), T.bfloat16
|
||||
)
|
||||
|
||||
# scores：FP32 QK 累加器和当前 tile 的指数权重。
|
||||
scores = T.alloc_fragment(
|
||||
(packed_block_m, packed_block_n), T.float32
|
||||
)
|
||||
# probs：scores 转成 BF16 后供 PV 矩阵乘使用。
|
||||
probs = T.alloc_fragment(
|
||||
(packed_block_m, packed_block_n), T.bfloat16
|
||||
)
|
||||
# output_accum：跨全部 KV tile 保存 FP32 输出分子。
|
||||
output_accum = T.alloc_fragment(
|
||||
(packed_block_m, head_dim_vo), T.float32
|
||||
)
|
||||
# 每个 packed row 的在线 Softmax 状态。row_max_prev 是 scratch，
|
||||
# row_scale 在循环中保存历史尺度修正，循环结束后复用为 1/denom。
|
||||
row_max = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_max_prev = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_scale = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_denom = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
|
||||
# dense equal-length 特化：直接用 batch_idx*seq_len 定位请求，并将
|
||||
# q_len/kv_len 变成编译期常量；ragged 路径仍从 indptr 读取真实边界。
|
||||
# 这可删除 dense 热路径上的 indptr load 和部分动态边界判断。
|
||||
q_start = (
|
||||
batch_idx * seq_len
|
||||
if dense_equal
|
||||
else qo_indptr[batch_idx]
|
||||
)
|
||||
kv_start = (
|
||||
batch_idx * seq_len
|
||||
if dense_equal
|
||||
else kv_indptr[batch_idx]
|
||||
)
|
||||
q_len = (
|
||||
seq_len
|
||||
if dense_equal
|
||||
else qo_indptr[batch_idx + 1] - q_start
|
||||
)
|
||||
kv_len = (
|
||||
seq_len
|
||||
if dense_equal
|
||||
else kv_indptr[batch_idx + 1] - kv_start
|
||||
)
|
||||
# MetaX PipelinePlanning 在 dense 长度完全常量化后，会把同一
|
||||
# kv_shared 中先写 K、后写 V 的合法生命周期复用误判成 stage
|
||||
# 重叠写。dense 路径仅为 copy 边界保留一次运行时 indptr 读取，
|
||||
# 数学边界、网格和有效 tile 判断仍然使用常量 kv_len。
|
||||
copy_kv_len = (
|
||||
kv_indptr[batch_idx + 1] - kv_start
|
||||
if dense_equal
|
||||
else kv_len
|
||||
)
|
||||
# packed_q_len 是当前请求包含的逻辑 packed rows 数量。
|
||||
packed_q_len = q_len * group_size
|
||||
# physical packed_tile 来自 blockIdx.x；logical_packed_tile 决定实际
|
||||
# 处理哪段 Q。反转只改变 CTA 提交顺序，不改变数学结果或输出位置。
|
||||
logical_packed_tile = (
|
||||
packed_tile_count - 1 - packed_tile
|
||||
if reverse_dense_tiles
|
||||
else packed_tile
|
||||
)
|
||||
packed_tile_start = logical_packed_tile * packed_block_m
|
||||
# dense 路径的网格恰好覆盖全部 packed rows，因此所有 CTA 有效；
|
||||
# ragged 请求可能短于 seq_len 上界，需要在运行时过滤无效 CTA。
|
||||
valid_q_tile = (
|
||||
True
|
||||
if dense_equal
|
||||
else packed_tile_start < packed_q_len
|
||||
)
|
||||
# bottom-right causal mask 使用 kv_len-q_len 修正 Q/KV 长度差。
|
||||
causal_offset = kv_len - q_len
|
||||
|
||||
if valid_q_tile:
|
||||
# 将 packed row 反解为 query position 和 GQA 组内 Q head：
|
||||
# q_pos = packed_row // group_size
|
||||
# group_head= packed_row % group_size
|
||||
# 再与当前 kv_head 组合出原始 qo_head。最后一个 tile 的补齐行写0。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_qk):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
group_head = packed_row % group_size
|
||||
if packed_row < packed_q_len:
|
||||
q_shared[i, d] = q[
|
||||
q_start + q_pos,
|
||||
kv_head * group_size + group_head,
|
||||
d,
|
||||
]
|
||||
else:
|
||||
q_shared[i, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
# 初始化在线 Softmax：输出分子=0、分母=0、running max=-inf。
|
||||
T.clear(output_accum)
|
||||
T.clear(row_denom)
|
||||
T.fill(row_max, -T.infinity(T.float32))
|
||||
|
||||
# 当前 packed tile 末尾对应的 query position 上界。一个 packed tile
|
||||
# 覆盖 packed_block_m/group_size 个不同 query position。
|
||||
q_upper_bound = T.min(
|
||||
q_len,
|
||||
T.ceildiv(
|
||||
(logical_packed_tile + 1) * packed_block_m,
|
||||
group_size,
|
||||
),
|
||||
)
|
||||
# 使用 tile 内最后一个 query 的因果边界求最大可见 KV 长度，以便
|
||||
# 整块跳过其右侧不可能被任何行看到的 KV tiles。
|
||||
max_visible = T.min(
|
||||
T.max(0, causal_offset + q_upper_bound),
|
||||
kv_len,
|
||||
)
|
||||
visible_tile_count = (
|
||||
T.ceildiv(max_visible, packed_block_n)
|
||||
if causal
|
||||
else T.ceildiv(kv_len, packed_block_n)
|
||||
)
|
||||
# 无效 ragged Q tile 的 loop_range=0，可跳过全部主要计算。
|
||||
loop_range = T.if_then_else(
|
||||
valid_q_tile, visible_tile_count, 0
|
||||
)
|
||||
# 当前 packed tile 第一行对应的 query position。
|
||||
first_q_pos = packed_tile_start // group_size
|
||||
# 如果某个 KV tile 连当前 packed tile 的第一行都完全可见，那么它
|
||||
# 对后续所有行也完全可见。这部分 tile 无需逐元素生成 causal mask，
|
||||
# 直接把 scores 清零作为 GEMM 初始累加器即可。
|
||||
fully_visible_tiles = T.min(
|
||||
loop_range,
|
||||
T.max(
|
||||
0,
|
||||
(first_q_pos + causal_offset + 1) // packed_block_n,
|
||||
),
|
||||
)
|
||||
|
||||
# 顺序扫描当前 packed Q tile 可见的 KV tiles。在线 Softmax 的
|
||||
# running max/denom/output_accum 在相邻迭代间存在严格依赖。
|
||||
for kv_tile in T.Pipelined(loop_range, num_stages=NUM_STAGES):
|
||||
# 当前 KV tile 在请求内部的半开区间 [tile_start, tile_end)。
|
||||
tile_start = kv_tile * packed_block_n
|
||||
tile_end = tile_start + packed_block_n
|
||||
|
||||
if tile_end <= copy_kv_len:
|
||||
# 完整 K tile 走向量化 copy。disable_tma=True 使用当前
|
||||
# MetaX 后端已验证可工作的普通 shared-memory copy 路径。
|
||||
T.copy(
|
||||
k[
|
||||
kv_start + tile_start : kv_start + tile_end,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
kv_shared,
|
||||
disable_tma=True,
|
||||
)
|
||||
else:
|
||||
# 最后一个 ragged K tile 可能越过请求边界，必须逐元素判断；
|
||||
# 越界位置填0，不能读到扁平张量中下一个请求的数据。
|
||||
for j, d in T.Parallel(packed_block_n, head_dim_qk):
|
||||
kv_pos = tile_start + j
|
||||
if kv_pos < copy_kv_len:
|
||||
kv_shared[j, d] = k[
|
||||
kv_start + kv_pos, kv_head, d
|
||||
]
|
||||
else:
|
||||
kv_shared[j, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
if causal and kv_tile < fully_visible_tiles:
|
||||
# 整块对所有 query 行可见，只需将 QK 累加器初始化为0。
|
||||
T.clear(scores)
|
||||
else:
|
||||
# causal frontier 或 ragged tail 需要逐元素初始化 mask。
|
||||
# 合法条件同时检查 packed Q tail、KV tail 和 bottom-right
|
||||
# causal 边界；非法元素设为 -inf，使其 softmax 权重为0。
|
||||
for i, j in T.Parallel(
|
||||
packed_block_m, packed_block_n
|
||||
):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
kv_pos = tile_start + j
|
||||
scores[i, j] = T.if_then_else(
|
||||
(packed_row < packed_q_len)
|
||||
and (kv_pos < kv_len)
|
||||
and (
|
||||
(not causal)
|
||||
or kv_pos < q_pos + 1 + causal_offset
|
||||
),
|
||||
0.0,
|
||||
-T.infinity(T.float32),
|
||||
)
|
||||
|
||||
# scores = mask + Q @ K^T。clear_accum 默认为 False，因此前面
|
||||
# 写入的 0/-inf mask 会作为矩阵乘累加器被保留下来。
|
||||
T.gemm(
|
||||
q_shared,
|
||||
kv_shared,
|
||||
scores,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
# 求当前 tile 的逐行最大值。row_max_prev 是 scratch，row_max
|
||||
# 始终保存此前全部 KV tiles 的 running max。
|
||||
T.reduce_max(scores, row_max_prev, dim=1)
|
||||
|
||||
# QK 和 max reduction 已完成，K 不再被任何线程读取。此时立即
|
||||
# 用 V 覆盖 kv_shared，再执行只依赖寄存器 fragment 的 softmax。
|
||||
# opt_012 把 V copy 放在 softmax 之后，生成代码会在 reduction
|
||||
# barrier 和 V copy 之间再插入一次 __syncthreads()。提前 copy
|
||||
# 让 K/V 生命周期切换贴近已有 reduction 同步，并让 V 在 PV
|
||||
# GEMM 前更早就绪；数学依赖和 shared footprint 都保持不变。
|
||||
if tile_end <= copy_kv_len:
|
||||
T.copy(
|
||||
v[
|
||||
kv_start + tile_start : kv_start + tile_end,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
kv_shared,
|
||||
disable_tma=True,
|
||||
)
|
||||
else:
|
||||
for j, d in T.Parallel(packed_block_n, head_dim_vo):
|
||||
kv_pos = tile_start + j
|
||||
if kv_pos < copy_kv_len:
|
||||
kv_shared[j, d] = v[
|
||||
kv_start + kv_pos, kv_head, d
|
||||
]
|
||||
else:
|
||||
kv_shared[j, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
for i in T.Parallel(packed_block_m):
|
||||
# 最大值改变时，用 exp(old_max-new_max) 把历史分母和历史
|
||||
# 输出分子转换到新的数值尺度，保证在线 Softmax 数值稳定。
|
||||
row_scale[i] = T.exp2(
|
||||
row_max[i] * softmax_scale
|
||||
- T.max(row_max[i], row_max_prev[i])
|
||||
* softmax_scale
|
||||
)
|
||||
# 合并历史最大值与当前 tile 最大值。
|
||||
row_max[i] = T.max(row_max[i], row_max_prev[i])
|
||||
|
||||
# 将 logits 转成相对于新 running max 的指数值。采用
|
||||
# FlashAttention 风格在线 Softmax，不保存完整 attention 矩阵。
|
||||
for i, j in T.Parallel(packed_block_m, packed_block_n):
|
||||
scores[i, j] = T.exp2(
|
||||
scores[i, j] * softmax_scale
|
||||
- row_max[i] * softmax_scale
|
||||
)
|
||||
# tile max 已无后续用途，复用 row_max_prev 保存当前指数和。
|
||||
T.reduce_sum(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(packed_block_m):
|
||||
# 分母更新：历史分母先乘 row_scale，再加当前 tile 指数和。
|
||||
row_denom[i] = (
|
||||
row_denom[i] * row_scale[i] + row_max_prev[i]
|
||||
)
|
||||
# PV GEMM 输入为 BF16，将 FP32 指数权重转换到 probs fragment。
|
||||
T.copy(scores, probs)
|
||||
|
||||
# 历史输出分子与分母使用相同的 row_scale 重新定标。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 累积当前 tile 的输出贡献：output_accum += probs @ V。
|
||||
T.gemm(
|
||||
probs,
|
||||
kv_shared,
|
||||
output_accum,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# 所有 KV tiles 处理完后再归一化。row_scale 已结束循环内使命，
|
||||
# 现在复用为每行的 1/softmax_denominator。每行只除一次，后续
|
||||
# 128 个输出维度都通过乘法归一化。补齐行和完全 mask 行设为0。
|
||||
for i in T.Parallel(packed_block_m):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
row_scale[i] = T.if_then_else(
|
||||
(packed_row < packed_q_len)
|
||||
and (
|
||||
(not causal)
|
||||
or q_pos + causal_offset >= 0
|
||||
),
|
||||
1.0 / row_denom[i],
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 将 packed row 重新映射到原始 output[q_pos, qo_head, d]。
|
||||
# packed tail 只参与内部补齐，不允许写回输出张量。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_vo):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
group_head = packed_row % group_size
|
||||
if packed_row < packed_q_len:
|
||||
output[
|
||||
q_start + q_pos,
|
||||
kv_head * group_size + group_head,
|
||||
d,
|
||||
] = output_accum[i, d] * row_scale[i]
|
||||
|
||||
return packed_kernel_v_prefetch_after_qk_v15
|
||||
|
||||
|
||||
def run_kernel(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""OJ 调用入口：按完整规格选择、缓存并启动 TileLang kernel。"""
|
||||
|
||||
# cache key 必须包含所有会改变生成代码、网格或张量 shape 的参数：
|
||||
# q/k 的总长度决定静态张量形状；seq_len 决定网格上界、tile 配置和分支；
|
||||
# head 数、head_dim、causal 则直接影响 GQA 映射和 attention 数学逻辑。
|
||||
key = (
|
||||
"packed-gqa-v-prefetch-after-qk-v15",
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
if key not in _kernel_cache:
|
||||
# 首次遇到该参数组合时才执行 JIT：
|
||||
# seq_len<=128：小序列回退路径，每 CTA 处理一个 Q head，启动和资源
|
||||
# 开销更小，保留单 token、非2次幂尾部的现有优势；
|
||||
# seq_len>128 ：主 packed GQA 路径，在一个 CTA 内处理同一 KV head
|
||||
# 对应的多个 Q heads，提高 K/V 数据复用。
|
||||
if seq_len <= 128:
|
||||
_kernel_cache[key] = build_kernel(
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
else:
|
||||
_kernel_cache[key] = build_packed_kernel_v_prefetch_after_qk_v15(
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
# JIT 结果是可直接接收 Torch tensors 的 Cython backend callable。
|
||||
# output 由调用方提前分配，本函数只启动 kernel，不在计时路径创建临时张量。
|
||||
_kernel_cache[key](q, k, v, output, qo_indptr, kv_indptr)
|
||||
|
|
@ -0,0 +1,724 @@
|
|||
import tilelang
|
||||
import tilelang.language as T
|
||||
from tilelang import jit
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 通用配置
|
||||
# ---------------------------------------------------------------------------
|
||||
# 小序列回退 kernel 每个 CTA 处理 BLOCK_M 个 query position。
|
||||
BLOCK_M = 64
|
||||
# seq_len <= 128 时使用的 KV tile 宽度。
|
||||
BLOCK_N = 64
|
||||
# 较长回退形状使用更窄的 KV tile，降低 fragment 和 shared memory 压力。
|
||||
GENERAL_BLOCK_N = 32
|
||||
# 小序列回退 kernel 的线程数；主 packed kernel 会在构建时选择 256 或 512。
|
||||
NUM_THREADS = 128
|
||||
# 当前只使用单阶段循环。K 和 V 在同一块 shared memory 中分阶段复用，
|
||||
# 不能在没有重新设计双缓冲的情况下直接把这里改成 2。
|
||||
NUM_STAGES = 1
|
||||
# softmax 数学形式使用 exp，但设备上的 exp2 指令通常更高效，因此将
|
||||
# exp(x) 转换为 exp2(x * log2(e))。
|
||||
LOG2_E = 1.44269504
|
||||
# TileLang 会针对完整参数组合生成专用 kernel。缓存用于避免同一进程中相同
|
||||
# shape 的重复调用再次触发 JIT 编译；编译时间不应进入 kernel 性能测量。
|
||||
_kernel_cache = {}
|
||||
|
||||
|
||||
@jit(
|
||||
execution_backend="cython",
|
||||
pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True},
|
||||
)
|
||||
def build_kernel(
|
||||
total_q,
|
||||
total_kv,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""构建小序列回退 kernel：每个 CTA 独立处理一个 Q head。"""
|
||||
|
||||
# GQA 中多个 Q head 共享一个 KV head。本题 32 个 Q head、4 个 KV head，
|
||||
# 因此 group_size=8，Q head h 对应的 KV head 为 h // 8。
|
||||
group_size = num_qo_heads // num_kv_heads
|
||||
# 标准 attention scale 为 1/sqrt(head_dim_qk)，再乘 log2(e) 供 exp2 使用。
|
||||
softmax_scale = (1.0 / head_dim_qk) ** 0.5 * LOG2_E
|
||||
# 极短序列使用 N=64，减少 KV 循环次数；较长回退形状使用 N=32，减少
|
||||
# scores/probs fragment 和 shared memory 的瞬时占用。
|
||||
block_n = BLOCK_N if seq_len <= 128 else GENERAL_BLOCK_N
|
||||
# seq_len 是每个 ragged 请求长度的上界。若 total_q=batch_size*seq_len，
|
||||
# 所有请求都只能恰好等于 seq_len，此时网格中不存在无效 Q tile，可以让
|
||||
# JIT 在编译期删除 valid_q_tile 对循环范围的保护逻辑。
|
||||
guard_invalid_tiles = total_q != batch_size * seq_len
|
||||
|
||||
@T.prim_func
|
||||
def kernel(
|
||||
q: T.Tensor((total_q, num_qo_heads, head_dim_qk), T.bfloat16),
|
||||
k: T.Tensor((total_kv, num_kv_heads, head_dim_qk), T.bfloat16),
|
||||
v: T.Tensor((total_kv, num_kv_heads, head_dim_vo), T.bfloat16),
|
||||
output: T.Tensor((total_q, num_qo_heads, head_dim_vo), T.bfloat16),
|
||||
qo_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
kv_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
):
|
||||
# 网格维度：
|
||||
# x = 一个请求内的 Q tile 编号；
|
||||
# y = Q head 编号；
|
||||
# z = batch/request 编号。
|
||||
# 因此一个 CTA 负责 (batch_idx, qo_head, q_tile) 的完整 attention。
|
||||
with T.Kernel(
|
||||
T.ceildiv(seq_len, BLOCK_M),
|
||||
num_qo_heads,
|
||||
batch_size,
|
||||
threads=NUM_THREADS,
|
||||
) as (q_tile, qo_head, batch_idx):
|
||||
# shared memory 保存本 CTA 重复使用的 Q/K/V tile。Q 在整个 KV
|
||||
# 循环中保持不变，K 和 V 则随 kv_tile 更新。
|
||||
q_shared = T.alloc_shared((BLOCK_M, head_dim_qk), T.bfloat16)
|
||||
k_shared = T.alloc_shared((block_n, head_dim_qk), T.bfloat16)
|
||||
v_shared = T.alloc_shared((block_n, head_dim_vo), T.bfloat16)
|
||||
|
||||
# fragment 通常映射到线程私有寄存器/矩阵累加器：
|
||||
# scores : 当前 Q tile 与当前 K tile 的 FP32 logits；
|
||||
# probs : 将 softmax 权重转成 BF16，作为 PV GEMM 的输入；
|
||||
# output_accum : 跨所有 KV tile 累积的 FP32 输出分子。
|
||||
scores = T.alloc_fragment((BLOCK_M, block_n), T.float32)
|
||||
probs = T.alloc_fragment((BLOCK_M, block_n), T.bfloat16)
|
||||
output_accum = T.alloc_fragment((BLOCK_M, head_dim_vo), T.float32)
|
||||
# 在线 Softmax 每一行只维护少量状态：
|
||||
# row_max : 截止当前 KV tile 的全局最大 logit；
|
||||
# row_max_prev : 临时 scratch，先存当前 tile 最大值，后存指数和；
|
||||
# row_scale : 最大值变化后，历史累加结果需要乘的缩放因子；
|
||||
# row_denom : 截止当前 KV tile 的 softmax 分母。
|
||||
row_max = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_max_prev = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_scale = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_denom = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
|
||||
# indptr 给出当前 ragged 请求在扁平 Q/K/V 张量中的区间。
|
||||
q_start = qo_indptr[batch_idx]
|
||||
q_end = qo_indptr[batch_idx + 1]
|
||||
kv_start = kv_indptr[batch_idx]
|
||||
kv_end = kv_indptr[batch_idx + 1]
|
||||
q_len = q_end - q_start
|
||||
kv_len = kv_end - kv_start
|
||||
# 根据 GQA 分组找到该 Q head 共享的 KV head。
|
||||
kv_head = qo_head // group_size
|
||||
# bottom-right causal 对齐的偏移量。可见条件为：
|
||||
# kv_pos < q_pos + 1 + (kv_len - q_len)。
|
||||
causal_offset = kv_len - q_len
|
||||
# 网格按全局 seq_len 上界启动，ragged 请求可能没有对应的 q_tile。
|
||||
valid_q_tile = (
|
||||
q_tile * BLOCK_M < q_len if guard_invalid_tiles else True
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# Q tile 在整个 KV 循环中都会复用，所以只在循环前加载一次。
|
||||
T.copy(
|
||||
q[
|
||||
q_start + q_tile * BLOCK_M : q_start + (q_tile + 1) * BLOCK_M,
|
||||
qo_head,
|
||||
:,
|
||||
],
|
||||
q_shared,
|
||||
)
|
||||
# 在线 Softmax 初始状态：输出分子和分母为 0，最大值为 -inf。
|
||||
T.fill(output_accum, 0)
|
||||
T.fill(row_denom, 0)
|
||||
T.fill(row_max, -T.infinity(T.float32))
|
||||
|
||||
# q_tile 中最后一行 query 最多能看到的 KV 长度。提前缩短 KV 循环，
|
||||
# 避免对因果边界右侧完全不可见的 KV tile 执行 GEMM。
|
||||
max_visible = T.min(
|
||||
T.max(0, causal_offset + (q_tile + 1) * BLOCK_M),
|
||||
kv_len,
|
||||
)
|
||||
visible_tile_count = (
|
||||
T.ceildiv(max_visible, block_n)
|
||||
if causal
|
||||
else T.ceildiv(kv_len, block_n)
|
||||
)
|
||||
# 对无效 q_tile 将循环次数设为 0，从而跳过 K/V copy、两个 GEMM 和
|
||||
# Softmax，而不仅仅是在最后禁止写回。
|
||||
loop_range = (
|
||||
T.if_then_else(valid_q_tile, visible_tile_count, 0)
|
||||
if guard_invalid_tiles
|
||||
else visible_tile_count
|
||||
)
|
||||
|
||||
# 逐块扫描当前 Q tile 可见的 KV 区域。NUM_STAGES=1 表示这里没有
|
||||
# 跨 kv_tile 的 K/V 双缓冲，所有在线 Softmax 状态都存在循环依赖。
|
||||
for kv_tile in T.Pipelined(loop_range, num_stages=NUM_STAGES):
|
||||
# 加载当前 K tile。回退路径主要处理短序列，边界安全由 TileLang
|
||||
# 对 copy 的合法范围处理以及后续显式 mask 共同保证。
|
||||
T.copy(
|
||||
k[
|
||||
kv_start + kv_tile * block_n : kv_start + (kv_tile + 1) * block_n,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
k_shared,
|
||||
)
|
||||
|
||||
# 先把 scores 初始化为 mask：合法元素为 0，非法元素为大负数。
|
||||
# 后续 QK GEMM 默认累加到 scores，因此最终得到 QK 或被 mask 的
|
||||
# 大负数，而不是让 GEMM 覆盖掉这里的因果/尾部 mask。
|
||||
for i, j in T.Parallel(BLOCK_M, block_n):
|
||||
q_pos = q_tile * BLOCK_M + i
|
||||
kv_pos = kv_tile * block_n + j
|
||||
scores[i, j] = T.if_then_else(
|
||||
(q_pos >= q_len)
|
||||
or (kv_pos >= kv_len)
|
||||
or (causal and kv_pos >= q_pos + 1 + causal_offset),
|
||||
-1.0e9,
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 计算 Q @ K^T。默认 clear_accum=False，会保留上面写入 scores
|
||||
# 的初始 mask 并将矩阵乘结果累加进去。
|
||||
T.gemm(
|
||||
q_shared,
|
||||
k_shared,
|
||||
scores,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
# 第一步：求当前 KV tile 每一行的最大值，暂存在 row_max_prev。
|
||||
# row_max 仍然保存此前所有 KV tile 的 running max。
|
||||
T.reduce_max(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
# 若当前 tile 提高了最大值，历史分母和历史输出分子都必须乘：
|
||||
# exp(old_max-new_max)。这里已换算为 exp2 域。
|
||||
row_scale[i] = T.exp2(
|
||||
row_max[i] * softmax_scale
|
||||
- T.max(row_max[i], row_max_prev[i])
|
||||
* softmax_scale
|
||||
)
|
||||
# 将 running max 更新为包含当前 tile 的新最大值。
|
||||
row_max[i] = T.max(row_max[i], row_max_prev[i])
|
||||
|
||||
# 第二步：计算当前 tile 相对于新 running max 的未归一化指数值。
|
||||
# 此处不立即除以分母，避免每个 KV tile 都执行完整归一化。
|
||||
for i, j in T.Parallel(BLOCK_M, block_n):
|
||||
scores[i, j] = T.exp2(
|
||||
scores[i, j] * softmax_scale
|
||||
- row_max[i] * softmax_scale
|
||||
)
|
||||
# row_max_prev 的 tile-max 已完成使命，现在复用同一个 fragment
|
||||
# 保存当前 tile 每一行的指数和，从而避免单独分配 row_sum。
|
||||
T.reduce_sum(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
# 更新在线 Softmax 分母：先把历史分母调整到新最大值尺度，
|
||||
# 再加上当前 tile 的指数和。
|
||||
row_denom[i] = (
|
||||
row_denom[i] * row_scale[i] + row_max_prev[i]
|
||||
)
|
||||
# PV GEMM 使用 BF16 输入，因此把 FP32 指数权重转换到 probs。
|
||||
T.copy(scores, probs)
|
||||
|
||||
# 输出分子和分母必须处于相同的最大值尺度。先缩放历史输出分子，
|
||||
# 再通过下面的 PV GEMM 加入当前 tile 的贡献。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 当前 K 已经不再使用，加载对应的 V tile。
|
||||
T.copy(
|
||||
v[
|
||||
kv_start + kv_tile * block_n : kv_start + (kv_tile + 1) * block_n,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
v_shared,
|
||||
)
|
||||
# output_accum += probs @ V，保持 FP32 累加精度。
|
||||
T.gemm(
|
||||
probs,
|
||||
v_shared,
|
||||
output_accum,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# KV 循环结束后才做最终归一化。每行只计算一次 1/denom，避免
|
||||
# head_dim_vo 个输出元素分别执行相同的除法。
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
q_pos = q_tile * BLOCK_M + i
|
||||
row_scale[i] = T.if_then_else(
|
||||
(q_pos < q_len)
|
||||
and (
|
||||
(not causal)
|
||||
or q_pos + causal_offset >= 0
|
||||
),
|
||||
1.0 / row_denom[i],
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 用乘法完成整行归一化。完全被 mask 的行乘 0。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 只写回请求实际存在的 query 行，抑制最后一个 Q tile 的尾部。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
if q_tile * BLOCK_M + i < q_len:
|
||||
output[
|
||||
q_start + q_tile * BLOCK_M + i,
|
||||
qo_head,
|
||||
d,
|
||||
] = output_accum[i, d]
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
@jit(
|
||||
execution_backend="cython",
|
||||
pass_configs={
|
||||
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
|
||||
tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
|
||||
},
|
||||
compile_flags=["-O3", "-DENABLE_BF16"],
|
||||
)
|
||||
def build_packed_kernel_square_warp_partition_v16(
|
||||
total_q,
|
||||
total_kv,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""构建主 GQA kernel，并在长路径平衡 M/N 方向的 wave 分区。"""
|
||||
|
||||
# equal-length 且 seq_len<=1024 时使用 M64/N32/256 threads。较小的资源
|
||||
# 占用可提高 CTA 驻留数；长序列和非对称形状保持 M128/N64/512 threads，
|
||||
# 以减少 Q tile 数量、KV 循环次数和在线 Softmax 更新次数。
|
||||
use_resident2 = total_q == total_kv and seq_len <= 1024
|
||||
packed_block_m = 64 if use_resident2 else 128
|
||||
packed_block_n = 32 if use_resident2 else 64
|
||||
packed_threads = 256 if use_resident2 else 512
|
||||
# FullRow 将长路径的 8 个 waves 全部分配到 M 维。case-4 生成代码因此
|
||||
# 在 PV 阶段从 V shared buffer 执行标量 BF16 读取。QK 与 PV 必须使用
|
||||
# 相同 policy 才能让 scores/probs fragment 布局一致；长路径同时改成
|
||||
# Square，测试更均衡的 M/N wave 分区。短 resident2 路径保留 FullRow。
|
||||
gemm_policy = (
|
||||
T.GemmWarpPolicy.FullRow
|
||||
if use_resident2
|
||||
else T.GemmWarpPolicy.Square
|
||||
)
|
||||
# 本题 group_size=8。一个 KV head 对应连续的 8 个 Q head。
|
||||
group_size = num_qo_heads // num_kv_heads
|
||||
# packed M 维中的每一行不是单纯的 query position，而是：
|
||||
# packed_row = q_pos * group_size + group_head。
|
||||
# 因此同一个 CTA 可以处理一个 KV head 对应的多个 Q head，并让这些 Q
|
||||
# head 在 CTA 内共享 K/V 的 global->shared 加载。
|
||||
packed_tile_count = T.ceildiv(seq_len * group_size, packed_block_m)
|
||||
# dense_equal 是编译期常量。由于每个 segment 长度都不超过 seq_len，而
|
||||
# 总长度恰好为 batch_size*seq_len，可以严格推出每个 segment 都等长。
|
||||
dense_equal = (
|
||||
total_q == batch_size * seq_len
|
||||
and total_kv == batch_size * seq_len
|
||||
)
|
||||
# 对 dense 长序列启用反向逻辑 tile 映射。因果 attention 中越靠后的 Q
|
||||
# tile 能看到越多 KV，计算量越大；优先提交重 tile 可以缩短最后少数 CTA
|
||||
# 造成的调度长尾。resident2 和 ragged 路径仍保持正向映射。
|
||||
reverse_dense_tiles = dense_equal and not use_resident2
|
||||
# attention scale 转为 exp2 所需的 base-2 缩放系数。
|
||||
softmax_scale = (1.0 / head_dim_qk) ** 0.5 * LOG2_E
|
||||
|
||||
@T.prim_func
|
||||
def packed_kernel_square_warp_partition_v16(
|
||||
q: T.Tensor((total_q, num_qo_heads, head_dim_qk), T.bfloat16),
|
||||
k: T.Tensor((total_kv, num_kv_heads, head_dim_qk), T.bfloat16),
|
||||
v: T.Tensor((total_kv, num_kv_heads, head_dim_vo), T.bfloat16),
|
||||
output: T.Tensor((total_q, num_qo_heads, head_dim_vo), T.bfloat16),
|
||||
qo_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
kv_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
):
|
||||
# 网格维度：
|
||||
# x = packed Q tile；
|
||||
# y = KV head；
|
||||
# z = batch/request。
|
||||
# 一个 CTA 因而处理一个请求、一个 KV head 和一段 packed Q rows。
|
||||
with T.Kernel(
|
||||
packed_tile_count,
|
||||
num_kv_heads,
|
||||
batch_size,
|
||||
threads=packed_threads,
|
||||
) as (packed_tile, kv_head, batch_idx):
|
||||
# Q tile 在整个 KV 循环中常驻 shared memory。长路径 M=128 时，
|
||||
# q_shared 大小为 128*128*2 = 32 KiB。
|
||||
q_shared = T.alloc_shared(
|
||||
(packed_block_m, head_dim_qk), T.bfloat16
|
||||
)
|
||||
# 同一个 kv_shared 分时保存 K 和 V：QK GEMM 完成后当前 K 生命周期
|
||||
# 已结束，随后 V 覆盖这块空间。这样长路径只需 16 KiB KV shared，
|
||||
# 而不是分别为 K、V 分配两份内存。
|
||||
kv_shared = T.alloc_shared(
|
||||
(packed_block_n, head_dim_qk), T.bfloat16
|
||||
)
|
||||
|
||||
# scores：FP32 QK 累加器和当前 tile 的指数权重。
|
||||
scores = T.alloc_fragment(
|
||||
(packed_block_m, packed_block_n), T.float32
|
||||
)
|
||||
# probs：scores 转成 BF16 后供 PV 矩阵乘使用。
|
||||
probs = T.alloc_fragment(
|
||||
(packed_block_m, packed_block_n), T.bfloat16
|
||||
)
|
||||
# output_accum：跨全部 KV tile 保存 FP32 输出分子。
|
||||
output_accum = T.alloc_fragment(
|
||||
(packed_block_m, head_dim_vo), T.float32
|
||||
)
|
||||
# 每个 packed row 的在线 Softmax 状态。row_max_prev 是 scratch，
|
||||
# row_scale 在循环中保存历史尺度修正，循环结束后复用为 1/denom。
|
||||
row_max = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_max_prev = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_scale = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_denom = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
|
||||
# dense equal-length 特化：直接用 batch_idx*seq_len 定位请求，并将
|
||||
# q_len/kv_len 变成编译期常量；ragged 路径仍从 indptr 读取真实边界。
|
||||
# 这可删除 dense 热路径上的 indptr load 和部分动态边界判断。
|
||||
q_start = (
|
||||
batch_idx * seq_len
|
||||
if dense_equal
|
||||
else qo_indptr[batch_idx]
|
||||
)
|
||||
kv_start = (
|
||||
batch_idx * seq_len
|
||||
if dense_equal
|
||||
else kv_indptr[batch_idx]
|
||||
)
|
||||
q_len = (
|
||||
seq_len
|
||||
if dense_equal
|
||||
else qo_indptr[batch_idx + 1] - q_start
|
||||
)
|
||||
kv_len = (
|
||||
seq_len
|
||||
if dense_equal
|
||||
else kv_indptr[batch_idx + 1] - kv_start
|
||||
)
|
||||
# MetaX PipelinePlanning 在 dense 长度完全常量化后，会把同一
|
||||
# kv_shared 中先写 K、后写 V 的合法生命周期复用误判成 stage
|
||||
# 重叠写。dense 路径仅为 copy 边界保留一次运行时 indptr 读取，
|
||||
# 数学边界、网格和有效 tile 判断仍然使用常量 kv_len。
|
||||
copy_kv_len = (
|
||||
kv_indptr[batch_idx + 1] - kv_start
|
||||
if dense_equal
|
||||
else kv_len
|
||||
)
|
||||
# packed_q_len 是当前请求包含的逻辑 packed rows 数量。
|
||||
packed_q_len = q_len * group_size
|
||||
# physical packed_tile 来自 blockIdx.x；logical_packed_tile 决定实际
|
||||
# 处理哪段 Q。反转只改变 CTA 提交顺序，不改变数学结果或输出位置。
|
||||
logical_packed_tile = (
|
||||
packed_tile_count - 1 - packed_tile
|
||||
if reverse_dense_tiles
|
||||
else packed_tile
|
||||
)
|
||||
packed_tile_start = logical_packed_tile * packed_block_m
|
||||
# dense 路径的网格恰好覆盖全部 packed rows，因此所有 CTA 有效；
|
||||
# ragged 请求可能短于 seq_len 上界，需要在运行时过滤无效 CTA。
|
||||
valid_q_tile = (
|
||||
True
|
||||
if dense_equal
|
||||
else packed_tile_start < packed_q_len
|
||||
)
|
||||
# bottom-right causal mask 使用 kv_len-q_len 修正 Q/KV 长度差。
|
||||
causal_offset = kv_len - q_len
|
||||
|
||||
if valid_q_tile:
|
||||
# 将 packed row 反解为 query position 和 GQA 组内 Q head：
|
||||
# q_pos = packed_row // group_size
|
||||
# group_head= packed_row % group_size
|
||||
# 再与当前 kv_head 组合出原始 qo_head。最后一个 tile 的补齐行写0。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_qk):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
group_head = packed_row % group_size
|
||||
if packed_row < packed_q_len:
|
||||
q_shared[i, d] = q[
|
||||
q_start + q_pos,
|
||||
kv_head * group_size + group_head,
|
||||
d,
|
||||
]
|
||||
else:
|
||||
q_shared[i, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
# 初始化在线 Softmax：输出分子=0、分母=0、running max=-inf。
|
||||
T.clear(output_accum)
|
||||
T.clear(row_denom)
|
||||
T.fill(row_max, -T.infinity(T.float32))
|
||||
|
||||
# 当前 packed tile 末尾对应的 query position 上界。一个 packed tile
|
||||
# 覆盖 packed_block_m/group_size 个不同 query position。
|
||||
q_upper_bound = T.min(
|
||||
q_len,
|
||||
T.ceildiv(
|
||||
(logical_packed_tile + 1) * packed_block_m,
|
||||
group_size,
|
||||
),
|
||||
)
|
||||
# 使用 tile 内最后一个 query 的因果边界求最大可见 KV 长度，以便
|
||||
# 整块跳过其右侧不可能被任何行看到的 KV tiles。
|
||||
max_visible = T.min(
|
||||
T.max(0, causal_offset + q_upper_bound),
|
||||
kv_len,
|
||||
)
|
||||
visible_tile_count = (
|
||||
T.ceildiv(max_visible, packed_block_n)
|
||||
if causal
|
||||
else T.ceildiv(kv_len, packed_block_n)
|
||||
)
|
||||
# 无效 ragged Q tile 的 loop_range=0，可跳过全部主要计算。
|
||||
loop_range = T.if_then_else(
|
||||
valid_q_tile, visible_tile_count, 0
|
||||
)
|
||||
# 当前 packed tile 第一行对应的 query position。
|
||||
first_q_pos = packed_tile_start // group_size
|
||||
# 如果某个 KV tile 连当前 packed tile 的第一行都完全可见，那么它
|
||||
# 对后续所有行也完全可见。这部分 tile 无需逐元素生成 causal mask，
|
||||
# 直接把 scores 清零作为 GEMM 初始累加器即可。
|
||||
fully_visible_tiles = T.min(
|
||||
loop_range,
|
||||
T.max(
|
||||
0,
|
||||
(first_q_pos + causal_offset + 1) // packed_block_n,
|
||||
),
|
||||
)
|
||||
|
||||
# 顺序扫描当前 packed Q tile 可见的 KV tiles。在线 Softmax 的
|
||||
# running max/denom/output_accum 在相邻迭代间存在严格依赖。
|
||||
for kv_tile in T.Pipelined(loop_range, num_stages=NUM_STAGES):
|
||||
# 当前 KV tile 在请求内部的半开区间 [tile_start, tile_end)。
|
||||
tile_start = kv_tile * packed_block_n
|
||||
tile_end = tile_start + packed_block_n
|
||||
|
||||
if tile_end <= copy_kv_len:
|
||||
# 完整 K tile 走向量化 copy。disable_tma=True 使用当前
|
||||
# MetaX 后端已验证可工作的普通 shared-memory copy 路径。
|
||||
T.copy(
|
||||
k[
|
||||
kv_start + tile_start : kv_start + tile_end,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
kv_shared,
|
||||
disable_tma=True,
|
||||
)
|
||||
else:
|
||||
# 最后一个 ragged K tile 可能越过请求边界，必须逐元素判断；
|
||||
# 越界位置填0，不能读到扁平张量中下一个请求的数据。
|
||||
for j, d in T.Parallel(packed_block_n, head_dim_qk):
|
||||
kv_pos = tile_start + j
|
||||
if kv_pos < copy_kv_len:
|
||||
kv_shared[j, d] = k[
|
||||
kv_start + kv_pos, kv_head, d
|
||||
]
|
||||
else:
|
||||
kv_shared[j, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
if causal and kv_tile < fully_visible_tiles:
|
||||
# 整块对所有 query 行可见，只需将 QK 累加器初始化为0。
|
||||
T.clear(scores)
|
||||
else:
|
||||
# causal frontier 或 ragged tail 需要逐元素初始化 mask。
|
||||
# 合法条件同时检查 packed Q tail、KV tail 和 bottom-right
|
||||
# causal 边界；非法元素设为 -inf，使其 softmax 权重为0。
|
||||
for i, j in T.Parallel(
|
||||
packed_block_m, packed_block_n
|
||||
):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
kv_pos = tile_start + j
|
||||
scores[i, j] = T.if_then_else(
|
||||
(packed_row < packed_q_len)
|
||||
and (kv_pos < kv_len)
|
||||
and (
|
||||
(not causal)
|
||||
or kv_pos < q_pos + 1 + causal_offset
|
||||
),
|
||||
0.0,
|
||||
-T.infinity(T.float32),
|
||||
)
|
||||
|
||||
# scores = mask + Q @ K^T。clear_accum 默认为 False，因此前面
|
||||
# 写入的 0/-inf mask 会作为矩阵乘累加器被保留下来。
|
||||
T.gemm(
|
||||
q_shared,
|
||||
kv_shared,
|
||||
scores,
|
||||
transpose_B=True,
|
||||
policy=gemm_policy,
|
||||
)
|
||||
|
||||
# 求当前 tile 的逐行最大值。row_max_prev 是 scratch，row_max
|
||||
# 始终保存此前全部 KV tiles 的 running max。
|
||||
T.reduce_max(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(packed_block_m):
|
||||
# 最大值改变时，用 exp(old_max-new_max) 把历史分母和历史
|
||||
# 输出分子转换到新的数值尺度，保证在线 Softmax 数值稳定。
|
||||
row_scale[i] = T.exp2(
|
||||
row_max[i] * softmax_scale
|
||||
- T.max(row_max[i], row_max_prev[i])
|
||||
* softmax_scale
|
||||
)
|
||||
# 合并历史最大值与当前 tile 最大值。
|
||||
row_max[i] = T.max(row_max[i], row_max_prev[i])
|
||||
|
||||
# 将 logits 转成相对于新 running max 的指数值。采用
|
||||
# FlashAttention 风格在线 Softmax，不保存完整 attention 矩阵。
|
||||
for i, j in T.Parallel(packed_block_m, packed_block_n):
|
||||
scores[i, j] = T.exp2(
|
||||
scores[i, j] * softmax_scale
|
||||
- row_max[i] * softmax_scale
|
||||
)
|
||||
# tile max 已无后续用途，复用 row_max_prev 保存当前指数和。
|
||||
T.reduce_sum(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(packed_block_m):
|
||||
# 分母更新：历史分母先乘 row_scale，再加当前 tile 指数和。
|
||||
row_denom[i] = (
|
||||
row_denom[i] * row_scale[i] + row_max_prev[i]
|
||||
)
|
||||
# PV GEMM 输入为 BF16，将 FP32 指数权重转换到 probs fragment。
|
||||
T.copy(scores, probs)
|
||||
|
||||
# 历史输出分子与分母使用相同的 row_scale 重新定标。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
if tile_end <= copy_kv_len:
|
||||
# QK 已结束，当前 K 不再存活；把对应 V tile 覆盖写入同一个
|
||||
# kv_shared，以减少每 CTA 的动态 shared memory 占用。
|
||||
T.copy(
|
||||
v[
|
||||
kv_start + tile_start : kv_start + tile_end,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
kv_shared,
|
||||
disable_tma=True,
|
||||
)
|
||||
else:
|
||||
# ragged V tail 与 K tail 一样执行显式边界判断和补0。
|
||||
for j, d in T.Parallel(packed_block_n, head_dim_vo):
|
||||
kv_pos = tile_start + j
|
||||
if kv_pos < copy_kv_len:
|
||||
kv_shared[j, d] = v[
|
||||
kv_start + kv_pos, kv_head, d
|
||||
]
|
||||
else:
|
||||
kv_shared[j, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
# 累积当前 tile 的输出贡献：output_accum += probs @ V。
|
||||
T.gemm(
|
||||
probs,
|
||||
kv_shared,
|
||||
output_accum,
|
||||
policy=gemm_policy,
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# 所有 KV tiles 处理完后再归一化。row_scale 已结束循环内使命，
|
||||
# 现在复用为每行的 1/softmax_denominator。每行只除一次，后续
|
||||
# 128 个输出维度都通过乘法归一化。补齐行和完全 mask 行设为0。
|
||||
for i in T.Parallel(packed_block_m):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
row_scale[i] = T.if_then_else(
|
||||
(packed_row < packed_q_len)
|
||||
and (
|
||||
(not causal)
|
||||
or q_pos + causal_offset >= 0
|
||||
),
|
||||
1.0 / row_denom[i],
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 将 packed row 重新映射到原始 output[q_pos, qo_head, d]。
|
||||
# packed tail 只参与内部补齐，不允许写回输出张量。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_vo):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
group_head = packed_row % group_size
|
||||
if packed_row < packed_q_len:
|
||||
output[
|
||||
q_start + q_pos,
|
||||
kv_head * group_size + group_head,
|
||||
d,
|
||||
] = output_accum[i, d] * row_scale[i]
|
||||
|
||||
return packed_kernel_square_warp_partition_v16
|
||||
|
||||
|
||||
def run_kernel(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""OJ 调用入口：按完整规格选择、缓存并启动 TileLang kernel。"""
|
||||
|
||||
# cache key 必须包含所有会改变生成代码、网格或张量 shape 的参数：
|
||||
# q/k 的总长度决定静态张量形状；seq_len 决定网格上界、tile 配置和分支；
|
||||
# head 数、head_dim、causal 则直接影响 GQA 映射和 attention 数学逻辑。
|
||||
key = (
|
||||
"packed-gqa-square-warp-partition-v16",
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
if key not in _kernel_cache:
|
||||
# 首次遇到该参数组合时才执行 JIT：
|
||||
# seq_len<=128：小序列回退路径，每 CTA 处理一个 Q head，启动和资源
|
||||
# 开销更小，保留单 token、非2次幂尾部的现有优势；
|
||||
# seq_len>128 ：主 packed GQA 路径，在一个 CTA 内处理同一 KV head
|
||||
# 对应的多个 Q heads，提高 K/V 数据复用。
|
||||
if seq_len <= 128:
|
||||
_kernel_cache[key] = build_kernel(
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
else:
|
||||
_kernel_cache[key] = build_packed_kernel_square_warp_partition_v16(
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
# JIT 结果是可直接接收 Torch tensors 的 Cython backend callable。
|
||||
# output 由调用方提前分配，本函数只启动 kernel，不在计时路径创建临时张量。
|
||||
_kernel_cache[key](q, k, v, output, qo_indptr, kv_indptr)
|
||||
|
|
@ -0,0 +1,719 @@
|
|||
import tilelang
|
||||
import tilelang.language as T
|
||||
from tilelang import jit
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 通用配置
|
||||
# ---------------------------------------------------------------------------
|
||||
# 小序列回退 kernel 每个 CTA 处理 BLOCK_M 个 query position。
|
||||
BLOCK_M = 64
|
||||
# seq_len <= 128 时使用的 KV tile 宽度。
|
||||
BLOCK_N = 64
|
||||
# 较长回退形状使用更窄的 KV tile，降低 fragment 和 shared memory 压力。
|
||||
GENERAL_BLOCK_N = 32
|
||||
# 小序列回退 kernel 的线程数；主 packed kernel 会在构建时选择 256 或 512。
|
||||
NUM_THREADS = 128
|
||||
# 当前只使用单阶段循环。K 和 V 在同一块 shared memory 中分阶段复用，
|
||||
# 不能在没有重新设计双缓冲的情况下直接把这里改成 2。
|
||||
NUM_STAGES = 1
|
||||
# softmax 数学形式使用 exp，但设备上的 exp2 指令通常更高效，因此将
|
||||
# exp(x) 转换为 exp2(x * log2(e))。
|
||||
LOG2_E = 1.44269504
|
||||
# TileLang 会针对完整参数组合生成专用 kernel。缓存用于避免同一进程中相同
|
||||
# shape 的重复调用再次触发 JIT 编译；编译时间不应进入 kernel 性能测量。
|
||||
_kernel_cache = {}
|
||||
|
||||
|
||||
@jit(
|
||||
execution_backend="cython",
|
||||
pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True},
|
||||
)
|
||||
def build_kernel(
|
||||
total_q,
|
||||
total_kv,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""构建小序列回退 kernel：每个 CTA 独立处理一个 Q head。"""
|
||||
|
||||
# GQA 中多个 Q head 共享一个 KV head。本题 32 个 Q head、4 个 KV head，
|
||||
# 因此 group_size=8，Q head h 对应的 KV head 为 h // 8。
|
||||
group_size = num_qo_heads // num_kv_heads
|
||||
# 标准 attention scale 为 1/sqrt(head_dim_qk)，再乘 log2(e) 供 exp2 使用。
|
||||
softmax_scale = (1.0 / head_dim_qk) ** 0.5 * LOG2_E
|
||||
# 极短序列使用 N=64，减少 KV 循环次数；较长回退形状使用 N=32，减少
|
||||
# scores/probs fragment 和 shared memory 的瞬时占用。
|
||||
block_n = BLOCK_N if seq_len <= 128 else GENERAL_BLOCK_N
|
||||
# seq_len 是每个 ragged 请求长度的上界。若 total_q=batch_size*seq_len，
|
||||
# 所有请求都只能恰好等于 seq_len，此时网格中不存在无效 Q tile，可以让
|
||||
# JIT 在编译期删除 valid_q_tile 对循环范围的保护逻辑。
|
||||
guard_invalid_tiles = total_q != batch_size * seq_len
|
||||
|
||||
@T.prim_func
|
||||
def kernel(
|
||||
q: T.Tensor((total_q, num_qo_heads, head_dim_qk), T.bfloat16),
|
||||
k: T.Tensor((total_kv, num_kv_heads, head_dim_qk), T.bfloat16),
|
||||
v: T.Tensor((total_kv, num_kv_heads, head_dim_vo), T.bfloat16),
|
||||
output: T.Tensor((total_q, num_qo_heads, head_dim_vo), T.bfloat16),
|
||||
qo_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
kv_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
):
|
||||
# 网格维度：
|
||||
# x = 一个请求内的 Q tile 编号；
|
||||
# y = Q head 编号；
|
||||
# z = batch/request 编号。
|
||||
# 因此一个 CTA 负责 (batch_idx, qo_head, q_tile) 的完整 attention。
|
||||
with T.Kernel(
|
||||
T.ceildiv(seq_len, BLOCK_M),
|
||||
num_qo_heads,
|
||||
batch_size,
|
||||
threads=NUM_THREADS,
|
||||
) as (q_tile, qo_head, batch_idx):
|
||||
# shared memory 保存本 CTA 重复使用的 Q/K/V tile。Q 在整个 KV
|
||||
# 循环中保持不变，K 和 V 则随 kv_tile 更新。
|
||||
q_shared = T.alloc_shared((BLOCK_M, head_dim_qk), T.bfloat16)
|
||||
k_shared = T.alloc_shared((block_n, head_dim_qk), T.bfloat16)
|
||||
v_shared = T.alloc_shared((block_n, head_dim_vo), T.bfloat16)
|
||||
|
||||
# fragment 通常映射到线程私有寄存器/矩阵累加器：
|
||||
# scores : 当前 Q tile 与当前 K tile 的 FP32 logits；
|
||||
# probs : 将 softmax 权重转成 BF16，作为 PV GEMM 的输入；
|
||||
# output_accum : 跨所有 KV tile 累积的 FP32 输出分子。
|
||||
scores = T.alloc_fragment((BLOCK_M, block_n), T.float32)
|
||||
probs = T.alloc_fragment((BLOCK_M, block_n), T.bfloat16)
|
||||
output_accum = T.alloc_fragment((BLOCK_M, head_dim_vo), T.float32)
|
||||
# 在线 Softmax 每一行只维护少量状态：
|
||||
# row_max : 截止当前 KV tile 的全局最大 logit；
|
||||
# row_max_prev : 临时 scratch，先存当前 tile 最大值，后存指数和；
|
||||
# row_scale : 最大值变化后，历史累加结果需要乘的缩放因子；
|
||||
# row_denom : 截止当前 KV tile 的 softmax 分母。
|
||||
row_max = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_max_prev = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_scale = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_denom = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
|
||||
# indptr 给出当前 ragged 请求在扁平 Q/K/V 张量中的区间。
|
||||
q_start = qo_indptr[batch_idx]
|
||||
q_end = qo_indptr[batch_idx + 1]
|
||||
kv_start = kv_indptr[batch_idx]
|
||||
kv_end = kv_indptr[batch_idx + 1]
|
||||
q_len = q_end - q_start
|
||||
kv_len = kv_end - kv_start
|
||||
# 根据 GQA 分组找到该 Q head 共享的 KV head。
|
||||
kv_head = qo_head // group_size
|
||||
# bottom-right causal 对齐的偏移量。可见条件为：
|
||||
# kv_pos < q_pos + 1 + (kv_len - q_len)。
|
||||
causal_offset = kv_len - q_len
|
||||
# 网格按全局 seq_len 上界启动，ragged 请求可能没有对应的 q_tile。
|
||||
valid_q_tile = (
|
||||
q_tile * BLOCK_M < q_len if guard_invalid_tiles else True
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# Q tile 在整个 KV 循环中都会复用，所以只在循环前加载一次。
|
||||
T.copy(
|
||||
q[
|
||||
q_start + q_tile * BLOCK_M : q_start + (q_tile + 1) * BLOCK_M,
|
||||
qo_head,
|
||||
:,
|
||||
],
|
||||
q_shared,
|
||||
)
|
||||
# 在线 Softmax 初始状态：输出分子和分母为 0，最大值为 -inf。
|
||||
T.fill(output_accum, 0)
|
||||
T.fill(row_denom, 0)
|
||||
T.fill(row_max, -T.infinity(T.float32))
|
||||
|
||||
# q_tile 中最后一行 query 最多能看到的 KV 长度。提前缩短 KV 循环，
|
||||
# 避免对因果边界右侧完全不可见的 KV tile 执行 GEMM。
|
||||
max_visible = T.min(
|
||||
T.max(0, causal_offset + (q_tile + 1) * BLOCK_M),
|
||||
kv_len,
|
||||
)
|
||||
visible_tile_count = (
|
||||
T.ceildiv(max_visible, block_n)
|
||||
if causal
|
||||
else T.ceildiv(kv_len, block_n)
|
||||
)
|
||||
# 对无效 q_tile 将循环次数设为 0，从而跳过 K/V copy、两个 GEMM 和
|
||||
# Softmax，而不仅仅是在最后禁止写回。
|
||||
loop_range = (
|
||||
T.if_then_else(valid_q_tile, visible_tile_count, 0)
|
||||
if guard_invalid_tiles
|
||||
else visible_tile_count
|
||||
)
|
||||
|
||||
# 逐块扫描当前 Q tile 可见的 KV 区域。NUM_STAGES=1 表示这里没有
|
||||
# 跨 kv_tile 的 K/V 双缓冲，所有在线 Softmax 状态都存在循环依赖。
|
||||
for kv_tile in T.Pipelined(loop_range, num_stages=NUM_STAGES):
|
||||
# 加载当前 K tile。回退路径主要处理短序列，边界安全由 TileLang
|
||||
# 对 copy 的合法范围处理以及后续显式 mask 共同保证。
|
||||
T.copy(
|
||||
k[
|
||||
kv_start + kv_tile * block_n : kv_start + (kv_tile + 1) * block_n,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
k_shared,
|
||||
)
|
||||
|
||||
# 先把 scores 初始化为 mask：合法元素为 0，非法元素为大负数。
|
||||
# 后续 QK GEMM 默认累加到 scores，因此最终得到 QK 或被 mask 的
|
||||
# 大负数，而不是让 GEMM 覆盖掉这里的因果/尾部 mask。
|
||||
for i, j in T.Parallel(BLOCK_M, block_n):
|
||||
q_pos = q_tile * BLOCK_M + i
|
||||
kv_pos = kv_tile * block_n + j
|
||||
scores[i, j] = T.if_then_else(
|
||||
(q_pos >= q_len)
|
||||
or (kv_pos >= kv_len)
|
||||
or (causal and kv_pos >= q_pos + 1 + causal_offset),
|
||||
-1.0e9,
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 计算 Q @ K^T。默认 clear_accum=False，会保留上面写入 scores
|
||||
# 的初始 mask 并将矩阵乘结果累加进去。
|
||||
T.gemm(
|
||||
q_shared,
|
||||
k_shared,
|
||||
scores,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
# 第一步：求当前 KV tile 每一行的最大值，暂存在 row_max_prev。
|
||||
# row_max 仍然保存此前所有 KV tile 的 running max。
|
||||
T.reduce_max(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
# 若当前 tile 提高了最大值，历史分母和历史输出分子都必须乘：
|
||||
# exp(old_max-new_max)。这里已换算为 exp2 域。
|
||||
row_scale[i] = T.exp2(
|
||||
row_max[i] * softmax_scale
|
||||
- T.max(row_max[i], row_max_prev[i])
|
||||
* softmax_scale
|
||||
)
|
||||
# 将 running max 更新为包含当前 tile 的新最大值。
|
||||
row_max[i] = T.max(row_max[i], row_max_prev[i])
|
||||
|
||||
# 第二步：计算当前 tile 相对于新 running max 的未归一化指数值。
|
||||
# 此处不立即除以分母，避免每个 KV tile 都执行完整归一化。
|
||||
for i, j in T.Parallel(BLOCK_M, block_n):
|
||||
scores[i, j] = T.exp2(
|
||||
scores[i, j] * softmax_scale
|
||||
- row_max[i] * softmax_scale
|
||||
)
|
||||
# row_max_prev 的 tile-max 已完成使命，现在复用同一个 fragment
|
||||
# 保存当前 tile 每一行的指数和，从而避免单独分配 row_sum。
|
||||
T.reduce_sum(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
# 更新在线 Softmax 分母：先把历史分母调整到新最大值尺度，
|
||||
# 再加上当前 tile 的指数和。
|
||||
row_denom[i] = (
|
||||
row_denom[i] * row_scale[i] + row_max_prev[i]
|
||||
)
|
||||
# PV GEMM 使用 BF16 输入，因此把 FP32 指数权重转换到 probs。
|
||||
T.copy(scores, probs)
|
||||
|
||||
# 输出分子和分母必须处于相同的最大值尺度。先缩放历史输出分子，
|
||||
# 再通过下面的 PV GEMM 加入当前 tile 的贡献。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 当前 K 已经不再使用，加载对应的 V tile。
|
||||
T.copy(
|
||||
v[
|
||||
kv_start + kv_tile * block_n : kv_start + (kv_tile + 1) * block_n,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
v_shared,
|
||||
)
|
||||
# output_accum += probs @ V，保持 FP32 累加精度。
|
||||
T.gemm(
|
||||
probs,
|
||||
v_shared,
|
||||
output_accum,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# KV 循环结束后才做最终归一化。每行只计算一次 1/denom，避免
|
||||
# head_dim_vo 个输出元素分别执行相同的除法。
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
q_pos = q_tile * BLOCK_M + i
|
||||
row_scale[i] = T.if_then_else(
|
||||
(q_pos < q_len)
|
||||
and (
|
||||
(not causal)
|
||||
or q_pos + causal_offset >= 0
|
||||
),
|
||||
1.0 / row_denom[i],
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 用乘法完成整行归一化。完全被 mask 的行乘 0。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 只写回请求实际存在的 query 行，抑制最后一个 Q tile 的尾部。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
if q_tile * BLOCK_M + i < q_len:
|
||||
output[
|
||||
q_start + q_tile * BLOCK_M + i,
|
||||
qo_head,
|
||||
d,
|
||||
] = output_accum[i, d]
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
@jit(
|
||||
execution_backend="cython",
|
||||
pass_configs={
|
||||
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
|
||||
tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
|
||||
},
|
||||
compile_flags=["-O3", "-DENABLE_BF16"],
|
||||
)
|
||||
def build_packed_kernel_skip_identity_rescale_v17(
|
||||
total_q,
|
||||
total_kv,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""构建主 GQA kernel，并跳过值为 1 的历史输出重缩放。"""
|
||||
|
||||
# equal-length 且 seq_len<=1024 时使用 M64/N32/256 threads。较小的资源
|
||||
# 占用可提高 CTA 驻留数；长序列和非对称形状保持 M128/N64/512 threads，
|
||||
# 以减少 Q tile 数量、KV 循环次数和在线 Softmax 更新次数。
|
||||
use_resident2 = total_q == total_kv and seq_len <= 1024
|
||||
packed_block_m = 64 if use_resident2 else 128
|
||||
packed_block_n = 32 if use_resident2 else 64
|
||||
packed_threads = 256 if use_resident2 else 512
|
||||
# 本题 group_size=8。一个 KV head 对应连续的 8 个 Q head。
|
||||
group_size = num_qo_heads // num_kv_heads
|
||||
# packed M 维中的每一行不是单纯的 query position，而是：
|
||||
# packed_row = q_pos * group_size + group_head。
|
||||
# 因此同一个 CTA 可以处理一个 KV head 对应的多个 Q head，并让这些 Q
|
||||
# head 在 CTA 内共享 K/V 的 global->shared 加载。
|
||||
packed_tile_count = T.ceildiv(seq_len * group_size, packed_block_m)
|
||||
# dense_equal 是编译期常量。由于每个 segment 长度都不超过 seq_len，而
|
||||
# 总长度恰好为 batch_size*seq_len，可以严格推出每个 segment 都等长。
|
||||
dense_equal = (
|
||||
total_q == batch_size * seq_len
|
||||
and total_kv == batch_size * seq_len
|
||||
)
|
||||
# 对 dense 长序列启用反向逻辑 tile 映射。因果 attention 中越靠后的 Q
|
||||
# tile 能看到越多 KV，计算量越大；优先提交重 tile 可以缩短最后少数 CTA
|
||||
# 造成的调度长尾。resident2 和 ragged 路径仍保持正向映射。
|
||||
reverse_dense_tiles = dense_equal and not use_resident2
|
||||
# attention scale 转为 exp2 所需的 base-2 缩放系数。
|
||||
softmax_scale = (1.0 / head_dim_qk) ** 0.5 * LOG2_E
|
||||
|
||||
@T.prim_func
|
||||
def packed_kernel_skip_identity_rescale_v17(
|
||||
q: T.Tensor((total_q, num_qo_heads, head_dim_qk), T.bfloat16),
|
||||
k: T.Tensor((total_kv, num_kv_heads, head_dim_qk), T.bfloat16),
|
||||
v: T.Tensor((total_kv, num_kv_heads, head_dim_vo), T.bfloat16),
|
||||
output: T.Tensor((total_q, num_qo_heads, head_dim_vo), T.bfloat16),
|
||||
qo_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
kv_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
):
|
||||
# 网格维度：
|
||||
# x = packed Q tile；
|
||||
# y = KV head；
|
||||
# z = batch/request。
|
||||
# 一个 CTA 因而处理一个请求、一个 KV head 和一段 packed Q rows。
|
||||
with T.Kernel(
|
||||
packed_tile_count,
|
||||
num_kv_heads,
|
||||
batch_size,
|
||||
threads=packed_threads,
|
||||
) as (packed_tile, kv_head, batch_idx):
|
||||
# Q tile 在整个 KV 循环中常驻 shared memory。长路径 M=128 时，
|
||||
# q_shared 大小为 128*128*2 = 32 KiB。
|
||||
q_shared = T.alloc_shared(
|
||||
(packed_block_m, head_dim_qk), T.bfloat16
|
||||
)
|
||||
# 同一个 kv_shared 分时保存 K 和 V：QK GEMM 完成后当前 K 生命周期
|
||||
# 已结束，随后 V 覆盖这块空间。这样长路径只需 16 KiB KV shared，
|
||||
# 而不是分别为 K、V 分配两份内存。
|
||||
kv_shared = T.alloc_shared(
|
||||
(packed_block_n, head_dim_qk), T.bfloat16
|
||||
)
|
||||
|
||||
# scores：FP32 QK 累加器和当前 tile 的指数权重。
|
||||
scores = T.alloc_fragment(
|
||||
(packed_block_m, packed_block_n), T.float32
|
||||
)
|
||||
# probs：scores 转成 BF16 后供 PV 矩阵乘使用。
|
||||
probs = T.alloc_fragment(
|
||||
(packed_block_m, packed_block_n), T.bfloat16
|
||||
)
|
||||
# output_accum：跨全部 KV tile 保存 FP32 输出分子。
|
||||
output_accum = T.alloc_fragment(
|
||||
(packed_block_m, head_dim_vo), T.float32
|
||||
)
|
||||
# 每个 packed row 的在线 Softmax 状态。row_max_prev 是 scratch，
|
||||
# row_scale 在循环中保存历史尺度修正，循环结束后复用为 1/denom。
|
||||
row_max = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_max_prev = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_scale = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_denom = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
|
||||
# dense equal-length 特化：直接用 batch_idx*seq_len 定位请求，并将
|
||||
# q_len/kv_len 变成编译期常量；ragged 路径仍从 indptr 读取真实边界。
|
||||
# 这可删除 dense 热路径上的 indptr load 和部分动态边界判断。
|
||||
q_start = (
|
||||
batch_idx * seq_len
|
||||
if dense_equal
|
||||
else qo_indptr[batch_idx]
|
||||
)
|
||||
kv_start = (
|
||||
batch_idx * seq_len
|
||||
if dense_equal
|
||||
else kv_indptr[batch_idx]
|
||||
)
|
||||
q_len = (
|
||||
seq_len
|
||||
if dense_equal
|
||||
else qo_indptr[batch_idx + 1] - q_start
|
||||
)
|
||||
kv_len = (
|
||||
seq_len
|
||||
if dense_equal
|
||||
else kv_indptr[batch_idx + 1] - kv_start
|
||||
)
|
||||
# MetaX PipelinePlanning 在 dense 长度完全常量化后，会把同一
|
||||
# kv_shared 中先写 K、后写 V 的合法生命周期复用误判成 stage
|
||||
# 重叠写。dense 路径仅为 copy 边界保留一次运行时 indptr 读取，
|
||||
# 数学边界、网格和有效 tile 判断仍然使用常量 kv_len。
|
||||
copy_kv_len = (
|
||||
kv_indptr[batch_idx + 1] - kv_start
|
||||
if dense_equal
|
||||
else kv_len
|
||||
)
|
||||
# packed_q_len 是当前请求包含的逻辑 packed rows 数量。
|
||||
packed_q_len = q_len * group_size
|
||||
# physical packed_tile 来自 blockIdx.x；logical_packed_tile 决定实际
|
||||
# 处理哪段 Q。反转只改变 CTA 提交顺序，不改变数学结果或输出位置。
|
||||
logical_packed_tile = (
|
||||
packed_tile_count - 1 - packed_tile
|
||||
if reverse_dense_tiles
|
||||
else packed_tile
|
||||
)
|
||||
packed_tile_start = logical_packed_tile * packed_block_m
|
||||
# dense 路径的网格恰好覆盖全部 packed rows，因此所有 CTA 有效；
|
||||
# ragged 请求可能短于 seq_len 上界，需要在运行时过滤无效 CTA。
|
||||
valid_q_tile = (
|
||||
True
|
||||
if dense_equal
|
||||
else packed_tile_start < packed_q_len
|
||||
)
|
||||
# bottom-right causal mask 使用 kv_len-q_len 修正 Q/KV 长度差。
|
||||
causal_offset = kv_len - q_len
|
||||
|
||||
if valid_q_tile:
|
||||
# 将 packed row 反解为 query position 和 GQA 组内 Q head：
|
||||
# q_pos = packed_row // group_size
|
||||
# group_head= packed_row % group_size
|
||||
# 再与当前 kv_head 组合出原始 qo_head。最后一个 tile 的补齐行写0。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_qk):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
group_head = packed_row % group_size
|
||||
if packed_row < packed_q_len:
|
||||
q_shared[i, d] = q[
|
||||
q_start + q_pos,
|
||||
kv_head * group_size + group_head,
|
||||
d,
|
||||
]
|
||||
else:
|
||||
q_shared[i, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
# 初始化在线 Softmax：输出分子=0、分母=0、running max=-inf。
|
||||
T.clear(output_accum)
|
||||
T.clear(row_denom)
|
||||
T.fill(row_max, -T.infinity(T.float32))
|
||||
|
||||
# 当前 packed tile 末尾对应的 query position 上界。一个 packed tile
|
||||
# 覆盖 packed_block_m/group_size 个不同 query position。
|
||||
q_upper_bound = T.min(
|
||||
q_len,
|
||||
T.ceildiv(
|
||||
(logical_packed_tile + 1) * packed_block_m,
|
||||
group_size,
|
||||
),
|
||||
)
|
||||
# 使用 tile 内最后一个 query 的因果边界求最大可见 KV 长度，以便
|
||||
# 整块跳过其右侧不可能被任何行看到的 KV tiles。
|
||||
max_visible = T.min(
|
||||
T.max(0, causal_offset + q_upper_bound),
|
||||
kv_len,
|
||||
)
|
||||
visible_tile_count = (
|
||||
T.ceildiv(max_visible, packed_block_n)
|
||||
if causal
|
||||
else T.ceildiv(kv_len, packed_block_n)
|
||||
)
|
||||
# 无效 ragged Q tile 的 loop_range=0，可跳过全部主要计算。
|
||||
loop_range = T.if_then_else(
|
||||
valid_q_tile, visible_tile_count, 0
|
||||
)
|
||||
# 当前 packed tile 第一行对应的 query position。
|
||||
first_q_pos = packed_tile_start // group_size
|
||||
# 如果某个 KV tile 连当前 packed tile 的第一行都完全可见，那么它
|
||||
# 对后续所有行也完全可见。这部分 tile 无需逐元素生成 causal mask，
|
||||
# 直接把 scores 清零作为 GEMM 初始累加器即可。
|
||||
fully_visible_tiles = T.min(
|
||||
loop_range,
|
||||
T.max(
|
||||
0,
|
||||
(first_q_pos + causal_offset + 1) // packed_block_n,
|
||||
),
|
||||
)
|
||||
|
||||
# 顺序扫描当前 packed Q tile 可见的 KV tiles。在线 Softmax 的
|
||||
# running max/denom/output_accum 在相邻迭代间存在严格依赖。
|
||||
for kv_tile in T.Pipelined(loop_range, num_stages=NUM_STAGES):
|
||||
# 当前 KV tile 在请求内部的半开区间 [tile_start, tile_end)。
|
||||
tile_start = kv_tile * packed_block_n
|
||||
tile_end = tile_start + packed_block_n
|
||||
|
||||
if tile_end <= copy_kv_len:
|
||||
# 完整 K tile 走向量化 copy。disable_tma=True 使用当前
|
||||
# MetaX 后端已验证可工作的普通 shared-memory copy 路径。
|
||||
T.copy(
|
||||
k[
|
||||
kv_start + tile_start : kv_start + tile_end,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
kv_shared,
|
||||
disable_tma=True,
|
||||
)
|
||||
else:
|
||||
# 最后一个 ragged K tile 可能越过请求边界，必须逐元素判断；
|
||||
# 越界位置填0，不能读到扁平张量中下一个请求的数据。
|
||||
for j, d in T.Parallel(packed_block_n, head_dim_qk):
|
||||
kv_pos = tile_start + j
|
||||
if kv_pos < copy_kv_len:
|
||||
kv_shared[j, d] = k[
|
||||
kv_start + kv_pos, kv_head, d
|
||||
]
|
||||
else:
|
||||
kv_shared[j, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
if causal and kv_tile < fully_visible_tiles:
|
||||
# 整块对所有 query 行可见，只需将 QK 累加器初始化为0。
|
||||
T.clear(scores)
|
||||
else:
|
||||
# causal frontier 或 ragged tail 需要逐元素初始化 mask。
|
||||
# 合法条件同时检查 packed Q tail、KV tail 和 bottom-right
|
||||
# causal 边界；非法元素设为 -inf，使其 softmax 权重为0。
|
||||
for i, j in T.Parallel(
|
||||
packed_block_m, packed_block_n
|
||||
):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
kv_pos = tile_start + j
|
||||
scores[i, j] = T.if_then_else(
|
||||
(packed_row < packed_q_len)
|
||||
and (kv_pos < kv_len)
|
||||
and (
|
||||
(not causal)
|
||||
or kv_pos < q_pos + 1 + causal_offset
|
||||
),
|
||||
0.0,
|
||||
-T.infinity(T.float32),
|
||||
)
|
||||
|
||||
# scores = mask + Q @ K^T。clear_accum 默认为 False，因此前面
|
||||
# 写入的 0/-inf mask 会作为矩阵乘累加器被保留下来。
|
||||
T.gemm(
|
||||
q_shared,
|
||||
kv_shared,
|
||||
scores,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
# 求当前 tile 的逐行最大值。row_max_prev 是 scratch，row_max
|
||||
# 始终保存此前全部 KV tiles 的 running max。
|
||||
T.reduce_max(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(packed_block_m):
|
||||
# 最大值改变时，用 exp(old_max-new_max) 把历史分母和历史
|
||||
# 输出分子转换到新的数值尺度，保证在线 Softmax 数值稳定。
|
||||
row_scale[i] = T.exp2(
|
||||
row_max[i] * softmax_scale
|
||||
- T.max(row_max[i], row_max_prev[i])
|
||||
* softmax_scale
|
||||
)
|
||||
# 合并历史最大值与当前 tile 最大值。
|
||||
row_max[i] = T.max(row_max[i], row_max_prev[i])
|
||||
|
||||
# 将 logits 转成相对于新 running max 的指数值。采用
|
||||
# FlashAttention 风格在线 Softmax，不保存完整 attention 矩阵。
|
||||
for i, j in T.Parallel(packed_block_m, packed_block_n):
|
||||
scores[i, j] = T.exp2(
|
||||
scores[i, j] * softmax_scale
|
||||
- row_max[i] * softmax_scale
|
||||
)
|
||||
# tile max 已无后续用途，复用 row_max_prev 保存当前指数和。
|
||||
T.reduce_sum(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(packed_block_m):
|
||||
# 分母更新：历史分母先乘 row_scale，再加当前 tile 指数和。
|
||||
row_denom[i] = (
|
||||
row_denom[i] * row_scale[i] + row_max_prev[i]
|
||||
)
|
||||
# PV GEMM 输入为 BF16，将 FP32 指数权重转换到 probs fragment。
|
||||
T.copy(scores, probs)
|
||||
|
||||
# 历史输出分子与分母使用相同的 row_scale 重新定标。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_vo):
|
||||
# 当当前 tile 没有提高 running max 时，exp2(0) 精确为1。
|
||||
# 长序列扫描到后段后这是常见路径；跳过整行 128 个恒等
|
||||
# 乘法，保留只有最大值改变时才需要的实际重缩放。
|
||||
if row_scale[i] != 1.0:
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
if tile_end <= copy_kv_len:
|
||||
# QK 已结束，当前 K 不再存活；把对应 V tile 覆盖写入同一个
|
||||
# kv_shared，以减少每 CTA 的动态 shared memory 占用。
|
||||
T.copy(
|
||||
v[
|
||||
kv_start + tile_start : kv_start + tile_end,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
kv_shared,
|
||||
disable_tma=True,
|
||||
)
|
||||
else:
|
||||
# ragged V tail 与 K tail 一样执行显式边界判断和补0。
|
||||
for j, d in T.Parallel(packed_block_n, head_dim_vo):
|
||||
kv_pos = tile_start + j
|
||||
if kv_pos < copy_kv_len:
|
||||
kv_shared[j, d] = v[
|
||||
kv_start + kv_pos, kv_head, d
|
||||
]
|
||||
else:
|
||||
kv_shared[j, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
# 累积当前 tile 的输出贡献：output_accum += probs @ V。
|
||||
T.gemm(
|
||||
probs,
|
||||
kv_shared,
|
||||
output_accum,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# 所有 KV tiles 处理完后再归一化。row_scale 已结束循环内使命，
|
||||
# 现在复用为每行的 1/softmax_denominator。每行只除一次，后续
|
||||
# 128 个输出维度都通过乘法归一化。补齐行和完全 mask 行设为0。
|
||||
for i in T.Parallel(packed_block_m):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
row_scale[i] = T.if_then_else(
|
||||
(packed_row < packed_q_len)
|
||||
and (
|
||||
(not causal)
|
||||
or q_pos + causal_offset >= 0
|
||||
),
|
||||
1.0 / row_denom[i],
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 将 packed row 重新映射到原始 output[q_pos, qo_head, d]。
|
||||
# packed tail 只参与内部补齐，不允许写回输出张量。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_vo):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
group_head = packed_row % group_size
|
||||
if packed_row < packed_q_len:
|
||||
output[
|
||||
q_start + q_pos,
|
||||
kv_head * group_size + group_head,
|
||||
d,
|
||||
] = output_accum[i, d] * row_scale[i]
|
||||
|
||||
return packed_kernel_skip_identity_rescale_v17
|
||||
|
||||
|
||||
def run_kernel(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""OJ 调用入口：按完整规格选择、缓存并启动 TileLang kernel。"""
|
||||
|
||||
# cache key 必须包含所有会改变生成代码、网格或张量 shape 的参数：
|
||||
# q/k 的总长度决定静态张量形状；seq_len 决定网格上界、tile 配置和分支；
|
||||
# head 数、head_dim、causal 则直接影响 GQA 映射和 attention 数学逻辑。
|
||||
key = (
|
||||
"packed-gqa-skip-identity-rescale-v17",
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
if key not in _kernel_cache:
|
||||
# 首次遇到该参数组合时才执行 JIT：
|
||||
# seq_len<=128：小序列回退路径，每 CTA 处理一个 Q head，启动和资源
|
||||
# 开销更小，保留单 token、非2次幂尾部的现有优势；
|
||||
# seq_len>128 ：主 packed GQA 路径，在一个 CTA 内处理同一 KV head
|
||||
# 对应的多个 Q heads，提高 K/V 数据复用。
|
||||
if seq_len <= 128:
|
||||
_kernel_cache[key] = build_kernel(
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
else:
|
||||
_kernel_cache[key] = build_packed_kernel_skip_identity_rescale_v17(
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
# JIT 结果是可直接接收 Torch tensors 的 Cython backend callable。
|
||||
# output 由调用方提前分配，本函数只启动 kernel，不在计时路径创建临时张量。
|
||||
_kernel_cache[key](q, k, v, output, qo_indptr, kv_indptr)
|
||||
|
|
@ -0,0 +1,716 @@
|
|||
import tilelang
|
||||
import tilelang.language as T
|
||||
from tilelang import jit
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 通用配置
|
||||
# ---------------------------------------------------------------------------
|
||||
# 小序列回退 kernel 每个 CTA 处理 BLOCK_M 个 query position。
|
||||
BLOCK_M = 64
|
||||
# seq_len <= 128 时使用的 KV tile 宽度。
|
||||
BLOCK_N = 64
|
||||
# 较长回退形状使用更窄的 KV tile，降低 fragment 和 shared memory 压力。
|
||||
GENERAL_BLOCK_N = 32
|
||||
# 小序列回退 kernel 的线程数；主 packed kernel 会在构建时选择 256 或 512。
|
||||
NUM_THREADS = 128
|
||||
# 当前只使用单阶段循环。K 和 V 在同一块 shared memory 中分阶段复用，
|
||||
# 不能在没有重新设计双缓冲的情况下直接把这里改成 2。
|
||||
NUM_STAGES = 1
|
||||
# softmax 数学形式使用 exp，但设备上的 exp2 指令通常更高效，因此将
|
||||
# exp(x) 转换为 exp2(x * log2(e))。
|
||||
LOG2_E = 1.44269504
|
||||
# TileLang 会针对完整参数组合生成专用 kernel。缓存用于避免同一进程中相同
|
||||
# shape 的重复调用再次触发 JIT 编译；编译时间不应进入 kernel 性能测量。
|
||||
_kernel_cache = {}
|
||||
|
||||
|
||||
@jit(
|
||||
execution_backend="cython",
|
||||
pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True},
|
||||
)
|
||||
def build_kernel(
|
||||
total_q,
|
||||
total_kv,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""构建小序列回退 kernel：每个 CTA 独立处理一个 Q head。"""
|
||||
|
||||
# GQA 中多个 Q head 共享一个 KV head。本题 32 个 Q head、4 个 KV head，
|
||||
# 因此 group_size=8，Q head h 对应的 KV head 为 h // 8。
|
||||
group_size = num_qo_heads // num_kv_heads
|
||||
# 标准 attention scale 为 1/sqrt(head_dim_qk)，再乘 log2(e) 供 exp2 使用。
|
||||
softmax_scale = (1.0 / head_dim_qk) ** 0.5 * LOG2_E
|
||||
# 极短序列使用 N=64，减少 KV 循环次数；较长回退形状使用 N=32，减少
|
||||
# scores/probs fragment 和 shared memory 的瞬时占用。
|
||||
block_n = BLOCK_N if seq_len <= 128 else GENERAL_BLOCK_N
|
||||
# seq_len 是每个 ragged 请求长度的上界。若 total_q=batch_size*seq_len，
|
||||
# 所有请求都只能恰好等于 seq_len，此时网格中不存在无效 Q tile，可以让
|
||||
# JIT 在编译期删除 valid_q_tile 对循环范围的保护逻辑。
|
||||
guard_invalid_tiles = total_q != batch_size * seq_len
|
||||
|
||||
@T.prim_func
|
||||
def kernel(
|
||||
q: T.Tensor((total_q, num_qo_heads, head_dim_qk), T.bfloat16),
|
||||
k: T.Tensor((total_kv, num_kv_heads, head_dim_qk), T.bfloat16),
|
||||
v: T.Tensor((total_kv, num_kv_heads, head_dim_vo), T.bfloat16),
|
||||
output: T.Tensor((total_q, num_qo_heads, head_dim_vo), T.bfloat16),
|
||||
qo_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
kv_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
):
|
||||
# 网格维度：
|
||||
# x = 一个请求内的 Q tile 编号；
|
||||
# y = Q head 编号；
|
||||
# z = batch/request 编号。
|
||||
# 因此一个 CTA 负责 (batch_idx, qo_head, q_tile) 的完整 attention。
|
||||
with T.Kernel(
|
||||
T.ceildiv(seq_len, BLOCK_M),
|
||||
num_qo_heads,
|
||||
batch_size,
|
||||
threads=NUM_THREADS,
|
||||
) as (q_tile, qo_head, batch_idx):
|
||||
# shared memory 保存本 CTA 重复使用的 Q/K/V tile。Q 在整个 KV
|
||||
# 循环中保持不变，K 和 V 则随 kv_tile 更新。
|
||||
q_shared = T.alloc_shared((BLOCK_M, head_dim_qk), T.bfloat16)
|
||||
k_shared = T.alloc_shared((block_n, head_dim_qk), T.bfloat16)
|
||||
v_shared = T.alloc_shared((block_n, head_dim_vo), T.bfloat16)
|
||||
|
||||
# fragment 通常映射到线程私有寄存器/矩阵累加器：
|
||||
# scores : 当前 Q tile 与当前 K tile 的 FP32 logits；
|
||||
# probs : 将 softmax 权重转成 BF16，作为 PV GEMM 的输入；
|
||||
# output_accum : 跨所有 KV tile 累积的 FP32 输出分子。
|
||||
scores = T.alloc_fragment((BLOCK_M, block_n), T.float32)
|
||||
probs = T.alloc_fragment((BLOCK_M, block_n), T.bfloat16)
|
||||
output_accum = T.alloc_fragment((BLOCK_M, head_dim_vo), T.float32)
|
||||
# 在线 Softmax 每一行只维护少量状态：
|
||||
# row_max : 截止当前 KV tile 的全局最大 logit；
|
||||
# row_max_prev : 临时 scratch，先存当前 tile 最大值，后存指数和；
|
||||
# row_scale : 最大值变化后，历史累加结果需要乘的缩放因子；
|
||||
# row_denom : 截止当前 KV tile 的 softmax 分母。
|
||||
row_max = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_max_prev = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_scale = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
row_denom = T.alloc_fragment((BLOCK_M,), T.float32)
|
||||
|
||||
# indptr 给出当前 ragged 请求在扁平 Q/K/V 张量中的区间。
|
||||
q_start = qo_indptr[batch_idx]
|
||||
q_end = qo_indptr[batch_idx + 1]
|
||||
kv_start = kv_indptr[batch_idx]
|
||||
kv_end = kv_indptr[batch_idx + 1]
|
||||
q_len = q_end - q_start
|
||||
kv_len = kv_end - kv_start
|
||||
# 根据 GQA 分组找到该 Q head 共享的 KV head。
|
||||
kv_head = qo_head // group_size
|
||||
# bottom-right causal 对齐的偏移量。可见条件为：
|
||||
# kv_pos < q_pos + 1 + (kv_len - q_len)。
|
||||
causal_offset = kv_len - q_len
|
||||
# 网格按全局 seq_len 上界启动，ragged 请求可能没有对应的 q_tile。
|
||||
valid_q_tile = (
|
||||
q_tile * BLOCK_M < q_len if guard_invalid_tiles else True
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# Q tile 在整个 KV 循环中都会复用，所以只在循环前加载一次。
|
||||
T.copy(
|
||||
q[
|
||||
q_start + q_tile * BLOCK_M : q_start + (q_tile + 1) * BLOCK_M,
|
||||
qo_head,
|
||||
:,
|
||||
],
|
||||
q_shared,
|
||||
)
|
||||
# 在线 Softmax 初始状态：输出分子和分母为 0，最大值为 -inf。
|
||||
T.fill(output_accum, 0)
|
||||
T.fill(row_denom, 0)
|
||||
T.fill(row_max, -T.infinity(T.float32))
|
||||
|
||||
# q_tile 中最后一行 query 最多能看到的 KV 长度。提前缩短 KV 循环，
|
||||
# 避免对因果边界右侧完全不可见的 KV tile 执行 GEMM。
|
||||
max_visible = T.min(
|
||||
T.max(0, causal_offset + (q_tile + 1) * BLOCK_M),
|
||||
kv_len,
|
||||
)
|
||||
visible_tile_count = (
|
||||
T.ceildiv(max_visible, block_n)
|
||||
if causal
|
||||
else T.ceildiv(kv_len, block_n)
|
||||
)
|
||||
# 对无效 q_tile 将循环次数设为 0，从而跳过 K/V copy、两个 GEMM 和
|
||||
# Softmax，而不仅仅是在最后禁止写回。
|
||||
loop_range = (
|
||||
T.if_then_else(valid_q_tile, visible_tile_count, 0)
|
||||
if guard_invalid_tiles
|
||||
else visible_tile_count
|
||||
)
|
||||
|
||||
# 逐块扫描当前 Q tile 可见的 KV 区域。NUM_STAGES=1 表示这里没有
|
||||
# 跨 kv_tile 的 K/V 双缓冲，所有在线 Softmax 状态都存在循环依赖。
|
||||
for kv_tile in T.Pipelined(loop_range, num_stages=NUM_STAGES):
|
||||
# 加载当前 K tile。回退路径主要处理短序列，边界安全由 TileLang
|
||||
# 对 copy 的合法范围处理以及后续显式 mask 共同保证。
|
||||
T.copy(
|
||||
k[
|
||||
kv_start + kv_tile * block_n : kv_start + (kv_tile + 1) * block_n,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
k_shared,
|
||||
)
|
||||
|
||||
# 先把 scores 初始化为 mask：合法元素为 0，非法元素为大负数。
|
||||
# 后续 QK GEMM 默认累加到 scores，因此最终得到 QK 或被 mask 的
|
||||
# 大负数，而不是让 GEMM 覆盖掉这里的因果/尾部 mask。
|
||||
for i, j in T.Parallel(BLOCK_M, block_n):
|
||||
q_pos = q_tile * BLOCK_M + i
|
||||
kv_pos = kv_tile * block_n + j
|
||||
scores[i, j] = T.if_then_else(
|
||||
(q_pos >= q_len)
|
||||
or (kv_pos >= kv_len)
|
||||
or (causal and kv_pos >= q_pos + 1 + causal_offset),
|
||||
-1.0e9,
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 计算 Q @ K^T。默认 clear_accum=False，会保留上面写入 scores
|
||||
# 的初始 mask 并将矩阵乘结果累加进去。
|
||||
T.gemm(
|
||||
q_shared,
|
||||
k_shared,
|
||||
scores,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
# 第一步：求当前 KV tile 每一行的最大值，暂存在 row_max_prev。
|
||||
# row_max 仍然保存此前所有 KV tile 的 running max。
|
||||
T.reduce_max(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
# 若当前 tile 提高了最大值，历史分母和历史输出分子都必须乘：
|
||||
# exp(old_max-new_max)。这里已换算为 exp2 域。
|
||||
row_scale[i] = T.exp2(
|
||||
row_max[i] * softmax_scale
|
||||
- T.max(row_max[i], row_max_prev[i])
|
||||
* softmax_scale
|
||||
)
|
||||
# 将 running max 更新为包含当前 tile 的新最大值。
|
||||
row_max[i] = T.max(row_max[i], row_max_prev[i])
|
||||
|
||||
# 第二步：计算当前 tile 相对于新 running max 的未归一化指数值。
|
||||
# 此处不立即除以分母，避免每个 KV tile 都执行完整归一化。
|
||||
for i, j in T.Parallel(BLOCK_M, block_n):
|
||||
scores[i, j] = T.exp2(
|
||||
scores[i, j] * softmax_scale
|
||||
- row_max[i] * softmax_scale
|
||||
)
|
||||
# row_max_prev 的 tile-max 已完成使命，现在复用同一个 fragment
|
||||
# 保存当前 tile 每一行的指数和，从而避免单独分配 row_sum。
|
||||
T.reduce_sum(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
# 更新在线 Softmax 分母：先把历史分母调整到新最大值尺度，
|
||||
# 再加上当前 tile 的指数和。
|
||||
row_denom[i] = (
|
||||
row_denom[i] * row_scale[i] + row_max_prev[i]
|
||||
)
|
||||
# PV GEMM 使用 BF16 输入，因此把 FP32 指数权重转换到 probs。
|
||||
T.copy(scores, probs)
|
||||
|
||||
# 输出分子和分母必须处于相同的最大值尺度。先缩放历史输出分子，
|
||||
# 再通过下面的 PV GEMM 加入当前 tile 的贡献。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 当前 K 已经不再使用，加载对应的 V tile。
|
||||
T.copy(
|
||||
v[
|
||||
kv_start + kv_tile * block_n : kv_start + (kv_tile + 1) * block_n,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
v_shared,
|
||||
)
|
||||
# output_accum += probs @ V，保持 FP32 累加精度。
|
||||
T.gemm(
|
||||
probs,
|
||||
v_shared,
|
||||
output_accum,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# KV 循环结束后才做最终归一化。每行只计算一次 1/denom，避免
|
||||
# head_dim_vo 个输出元素分别执行相同的除法。
|
||||
for i in T.Parallel(BLOCK_M):
|
||||
q_pos = q_tile * BLOCK_M + i
|
||||
row_scale[i] = T.if_then_else(
|
||||
(q_pos < q_len)
|
||||
and (
|
||||
(not causal)
|
||||
or q_pos + causal_offset >= 0
|
||||
),
|
||||
1.0 / row_denom[i],
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 用乘法完成整行归一化。完全被 mask 的行乘 0。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
# 只写回请求实际存在的 query 行，抑制最后一个 Q tile 的尾部。
|
||||
for i, d in T.Parallel(BLOCK_M, head_dim_vo):
|
||||
if q_tile * BLOCK_M + i < q_len:
|
||||
output[
|
||||
q_start + q_tile * BLOCK_M + i,
|
||||
qo_head,
|
||||
d,
|
||||
] = output_accum[i, d]
|
||||
|
||||
return kernel
|
||||
|
||||
|
||||
@jit(
|
||||
execution_backend="cython",
|
||||
pass_configs={
|
||||
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
|
||||
tilelang.PassConfigKey.TL_DISABLE_DATA_RACE_CHECK: True,
|
||||
},
|
||||
compile_flags=["-O3", "-DENABLE_BF16"],
|
||||
)
|
||||
def build_packed_kernel_factor_softmax_scale_v18(
|
||||
total_q,
|
||||
total_kv,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""构建主 GQA kernel，并因式分解 softmax 的公共缩放系数。"""
|
||||
|
||||
# equal-length 且 seq_len<=1024 时使用 M64/N32/256 threads。较小的资源
|
||||
# 占用可提高 CTA 驻留数；长序列和非对称形状保持 M128/N64/512 threads，
|
||||
# 以减少 Q tile 数量、KV 循环次数和在线 Softmax 更新次数。
|
||||
use_resident2 = total_q == total_kv and seq_len <= 1024
|
||||
packed_block_m = 64 if use_resident2 else 128
|
||||
packed_block_n = 32 if use_resident2 else 64
|
||||
packed_threads = 256 if use_resident2 else 512
|
||||
# 本题 group_size=8。一个 KV head 对应连续的 8 个 Q head。
|
||||
group_size = num_qo_heads // num_kv_heads
|
||||
# packed M 维中的每一行不是单纯的 query position，而是：
|
||||
# packed_row = q_pos * group_size + group_head。
|
||||
# 因此同一个 CTA 可以处理一个 KV head 对应的多个 Q head，并让这些 Q
|
||||
# head 在 CTA 内共享 K/V 的 global->shared 加载。
|
||||
packed_tile_count = T.ceildiv(seq_len * group_size, packed_block_m)
|
||||
# dense_equal 是编译期常量。由于每个 segment 长度都不超过 seq_len，而
|
||||
# 总长度恰好为 batch_size*seq_len，可以严格推出每个 segment 都等长。
|
||||
dense_equal = (
|
||||
total_q == batch_size * seq_len
|
||||
and total_kv == batch_size * seq_len
|
||||
)
|
||||
# 对 dense 长序列启用反向逻辑 tile 映射。因果 attention 中越靠后的 Q
|
||||
# tile 能看到越多 KV，计算量越大；优先提交重 tile 可以缩短最后少数 CTA
|
||||
# 造成的调度长尾。resident2 和 ragged 路径仍保持正向映射。
|
||||
reverse_dense_tiles = dense_equal and not use_resident2
|
||||
# attention scale 转为 exp2 所需的 base-2 缩放系数。
|
||||
softmax_scale = (1.0 / head_dim_qk) ** 0.5 * LOG2_E
|
||||
|
||||
@T.prim_func
|
||||
def packed_kernel_factor_softmax_scale_v18(
|
||||
q: T.Tensor((total_q, num_qo_heads, head_dim_qk), T.bfloat16),
|
||||
k: T.Tensor((total_kv, num_kv_heads, head_dim_qk), T.bfloat16),
|
||||
v: T.Tensor((total_kv, num_kv_heads, head_dim_vo), T.bfloat16),
|
||||
output: T.Tensor((total_q, num_qo_heads, head_dim_vo), T.bfloat16),
|
||||
qo_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
kv_indptr: T.Tensor((batch_size + 1,), T.int32),
|
||||
):
|
||||
# 网格维度：
|
||||
# x = packed Q tile；
|
||||
# y = KV head；
|
||||
# z = batch/request。
|
||||
# 一个 CTA 因而处理一个请求、一个 KV head 和一段 packed Q rows。
|
||||
with T.Kernel(
|
||||
packed_tile_count,
|
||||
num_kv_heads,
|
||||
batch_size,
|
||||
threads=packed_threads,
|
||||
) as (packed_tile, kv_head, batch_idx):
|
||||
# Q tile 在整个 KV 循环中常驻 shared memory。长路径 M=128 时，
|
||||
# q_shared 大小为 128*128*2 = 32 KiB。
|
||||
q_shared = T.alloc_shared(
|
||||
(packed_block_m, head_dim_qk), T.bfloat16
|
||||
)
|
||||
# 同一个 kv_shared 分时保存 K 和 V：QK GEMM 完成后当前 K 生命周期
|
||||
# 已结束，随后 V 覆盖这块空间。这样长路径只需 16 KiB KV shared，
|
||||
# 而不是分别为 K、V 分配两份内存。
|
||||
kv_shared = T.alloc_shared(
|
||||
(packed_block_n, head_dim_qk), T.bfloat16
|
||||
)
|
||||
|
||||
# scores：FP32 QK 累加器和当前 tile 的指数权重。
|
||||
scores = T.alloc_fragment(
|
||||
(packed_block_m, packed_block_n), T.float32
|
||||
)
|
||||
# probs：scores 转成 BF16 后供 PV 矩阵乘使用。
|
||||
probs = T.alloc_fragment(
|
||||
(packed_block_m, packed_block_n), T.bfloat16
|
||||
)
|
||||
# output_accum：跨全部 KV tile 保存 FP32 输出分子。
|
||||
output_accum = T.alloc_fragment(
|
||||
(packed_block_m, head_dim_vo), T.float32
|
||||
)
|
||||
# 每个 packed row 的在线 Softmax 状态。row_max_prev 是 scratch，
|
||||
# row_scale 在循环中保存历史尺度修正，循环结束后复用为 1/denom。
|
||||
row_max = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_max_prev = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_scale = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
row_denom = T.alloc_fragment((packed_block_m,), T.float32)
|
||||
|
||||
# dense equal-length 特化：直接用 batch_idx*seq_len 定位请求，并将
|
||||
# q_len/kv_len 变成编译期常量；ragged 路径仍从 indptr 读取真实边界。
|
||||
# 这可删除 dense 热路径上的 indptr load 和部分动态边界判断。
|
||||
q_start = (
|
||||
batch_idx * seq_len
|
||||
if dense_equal
|
||||
else qo_indptr[batch_idx]
|
||||
)
|
||||
kv_start = (
|
||||
batch_idx * seq_len
|
||||
if dense_equal
|
||||
else kv_indptr[batch_idx]
|
||||
)
|
||||
q_len = (
|
||||
seq_len
|
||||
if dense_equal
|
||||
else qo_indptr[batch_idx + 1] - q_start
|
||||
)
|
||||
kv_len = (
|
||||
seq_len
|
||||
if dense_equal
|
||||
else kv_indptr[batch_idx + 1] - kv_start
|
||||
)
|
||||
# MetaX PipelinePlanning 在 dense 长度完全常量化后，会把同一
|
||||
# kv_shared 中先写 K、后写 V 的合法生命周期复用误判成 stage
|
||||
# 重叠写。dense 路径仅为 copy 边界保留一次运行时 indptr 读取，
|
||||
# 数学边界、网格和有效 tile 判断仍然使用常量 kv_len。
|
||||
copy_kv_len = (
|
||||
kv_indptr[batch_idx + 1] - kv_start
|
||||
if dense_equal
|
||||
else kv_len
|
||||
)
|
||||
# packed_q_len 是当前请求包含的逻辑 packed rows 数量。
|
||||
packed_q_len = q_len * group_size
|
||||
# physical packed_tile 来自 blockIdx.x；logical_packed_tile 决定实际
|
||||
# 处理哪段 Q。反转只改变 CTA 提交顺序，不改变数学结果或输出位置。
|
||||
logical_packed_tile = (
|
||||
packed_tile_count - 1 - packed_tile
|
||||
if reverse_dense_tiles
|
||||
else packed_tile
|
||||
)
|
||||
packed_tile_start = logical_packed_tile * packed_block_m
|
||||
# dense 路径的网格恰好覆盖全部 packed rows，因此所有 CTA 有效；
|
||||
# ragged 请求可能短于 seq_len 上界，需要在运行时过滤无效 CTA。
|
||||
valid_q_tile = (
|
||||
True
|
||||
if dense_equal
|
||||
else packed_tile_start < packed_q_len
|
||||
)
|
||||
# bottom-right causal mask 使用 kv_len-q_len 修正 Q/KV 长度差。
|
||||
causal_offset = kv_len - q_len
|
||||
|
||||
if valid_q_tile:
|
||||
# 将 packed row 反解为 query position 和 GQA 组内 Q head：
|
||||
# q_pos = packed_row // group_size
|
||||
# group_head= packed_row % group_size
|
||||
# 再与当前 kv_head 组合出原始 qo_head。最后一个 tile 的补齐行写0。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_qk):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
group_head = packed_row % group_size
|
||||
if packed_row < packed_q_len:
|
||||
q_shared[i, d] = q[
|
||||
q_start + q_pos,
|
||||
kv_head * group_size + group_head,
|
||||
d,
|
||||
]
|
||||
else:
|
||||
q_shared[i, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
# 初始化在线 Softmax：输出分子=0、分母=0、running max=-inf。
|
||||
T.clear(output_accum)
|
||||
T.clear(row_denom)
|
||||
T.fill(row_max, -T.infinity(T.float32))
|
||||
|
||||
# 当前 packed tile 末尾对应的 query position 上界。一个 packed tile
|
||||
# 覆盖 packed_block_m/group_size 个不同 query position。
|
||||
q_upper_bound = T.min(
|
||||
q_len,
|
||||
T.ceildiv(
|
||||
(logical_packed_tile + 1) * packed_block_m,
|
||||
group_size,
|
||||
),
|
||||
)
|
||||
# 使用 tile 内最后一个 query 的因果边界求最大可见 KV 长度，以便
|
||||
# 整块跳过其右侧不可能被任何行看到的 KV tiles。
|
||||
max_visible = T.min(
|
||||
T.max(0, causal_offset + q_upper_bound),
|
||||
kv_len,
|
||||
)
|
||||
visible_tile_count = (
|
||||
T.ceildiv(max_visible, packed_block_n)
|
||||
if causal
|
||||
else T.ceildiv(kv_len, packed_block_n)
|
||||
)
|
||||
# 无效 ragged Q tile 的 loop_range=0，可跳过全部主要计算。
|
||||
loop_range = T.if_then_else(
|
||||
valid_q_tile, visible_tile_count, 0
|
||||
)
|
||||
# 当前 packed tile 第一行对应的 query position。
|
||||
first_q_pos = packed_tile_start // group_size
|
||||
# 如果某个 KV tile 连当前 packed tile 的第一行都完全可见，那么它
|
||||
# 对后续所有行也完全可见。这部分 tile 无需逐元素生成 causal mask，
|
||||
# 直接把 scores 清零作为 GEMM 初始累加器即可。
|
||||
fully_visible_tiles = T.min(
|
||||
loop_range,
|
||||
T.max(
|
||||
0,
|
||||
(first_q_pos + causal_offset + 1) // packed_block_n,
|
||||
),
|
||||
)
|
||||
|
||||
# 顺序扫描当前 packed Q tile 可见的 KV tiles。在线 Softmax 的
|
||||
# running max/denom/output_accum 在相邻迭代间存在严格依赖。
|
||||
for kv_tile in T.Pipelined(loop_range, num_stages=NUM_STAGES):
|
||||
# 当前 KV tile 在请求内部的半开区间 [tile_start, tile_end)。
|
||||
tile_start = kv_tile * packed_block_n
|
||||
tile_end = tile_start + packed_block_n
|
||||
|
||||
if tile_end <= copy_kv_len:
|
||||
# 完整 K tile 走向量化 copy。disable_tma=True 使用当前
|
||||
# MetaX 后端已验证可工作的普通 shared-memory copy 路径。
|
||||
T.copy(
|
||||
k[
|
||||
kv_start + tile_start : kv_start + tile_end,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
kv_shared,
|
||||
disable_tma=True,
|
||||
)
|
||||
else:
|
||||
# 最后一个 ragged K tile 可能越过请求边界，必须逐元素判断；
|
||||
# 越界位置填0，不能读到扁平张量中下一个请求的数据。
|
||||
for j, d in T.Parallel(packed_block_n, head_dim_qk):
|
||||
kv_pos = tile_start + j
|
||||
if kv_pos < copy_kv_len:
|
||||
kv_shared[j, d] = k[
|
||||
kv_start + kv_pos, kv_head, d
|
||||
]
|
||||
else:
|
||||
kv_shared[j, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
if causal and kv_tile < fully_visible_tiles:
|
||||
# 整块对所有 query 行可见，只需将 QK 累加器初始化为0。
|
||||
T.clear(scores)
|
||||
else:
|
||||
# causal frontier 或 ragged tail 需要逐元素初始化 mask。
|
||||
# 合法条件同时检查 packed Q tail、KV tail 和 bottom-right
|
||||
# causal 边界；非法元素设为 -inf，使其 softmax 权重为0。
|
||||
for i, j in T.Parallel(
|
||||
packed_block_m, packed_block_n
|
||||
):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
kv_pos = tile_start + j
|
||||
scores[i, j] = T.if_then_else(
|
||||
(packed_row < packed_q_len)
|
||||
and (kv_pos < kv_len)
|
||||
and (
|
||||
(not causal)
|
||||
or kv_pos < q_pos + 1 + causal_offset
|
||||
),
|
||||
0.0,
|
||||
-T.infinity(T.float32),
|
||||
)
|
||||
|
||||
# scores = mask + Q @ K^T。clear_accum 默认为 False，因此前面
|
||||
# 写入的 0/-inf mask 会作为矩阵乘累加器被保留下来。
|
||||
T.gemm(
|
||||
q_shared,
|
||||
kv_shared,
|
||||
scores,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
# 求当前 tile 的逐行最大值。row_max_prev 是 scratch，row_max
|
||||
# 始终保存此前全部 KV tiles 的 running max。
|
||||
T.reduce_max(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(packed_block_m):
|
||||
# 最大值改变时，用 exp(old_max-new_max) 把历史分母和历史
|
||||
# 输出分子转换到新的数值尺度，保证在线 Softmax 数值稳定。
|
||||
row_scale[i] = T.exp2(
|
||||
(
|
||||
row_max[i]
|
||||
- T.max(row_max[i], row_max_prev[i])
|
||||
)
|
||||
* softmax_scale
|
||||
)
|
||||
# 合并历史最大值与当前 tile 最大值。
|
||||
row_max[i] = T.max(row_max[i], row_max_prev[i])
|
||||
|
||||
# 将 logits 转成相对于新 running max 的指数值。采用
|
||||
# FlashAttention 风格在线 Softmax，不保存完整 attention 矩阵。
|
||||
for i, j in T.Parallel(packed_block_m, packed_block_n):
|
||||
scores[i, j] = T.exp2(
|
||||
(scores[i, j] - row_max[i]) * softmax_scale
|
||||
)
|
||||
# tile max 已无后续用途，复用 row_max_prev 保存当前指数和。
|
||||
T.reduce_sum(scores, row_max_prev, dim=1)
|
||||
for i in T.Parallel(packed_block_m):
|
||||
# 分母更新：历史分母先乘 row_scale，再加当前 tile 指数和。
|
||||
row_denom[i] = (
|
||||
row_denom[i] * row_scale[i] + row_max_prev[i]
|
||||
)
|
||||
# PV GEMM 输入为 BF16，将 FP32 指数权重转换到 probs fragment。
|
||||
T.copy(scores, probs)
|
||||
|
||||
# 历史输出分子与分母使用相同的 row_scale 重新定标。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_vo):
|
||||
output_accum[i, d] *= row_scale[i]
|
||||
|
||||
if tile_end <= copy_kv_len:
|
||||
# QK 已结束，当前 K 不再存活；把对应 V tile 覆盖写入同一个
|
||||
# kv_shared，以减少每 CTA 的动态 shared memory 占用。
|
||||
T.copy(
|
||||
v[
|
||||
kv_start + tile_start : kv_start + tile_end,
|
||||
kv_head,
|
||||
:,
|
||||
],
|
||||
kv_shared,
|
||||
disable_tma=True,
|
||||
)
|
||||
else:
|
||||
# ragged V tail 与 K tail 一样执行显式边界判断和补0。
|
||||
for j, d in T.Parallel(packed_block_n, head_dim_vo):
|
||||
kv_pos = tile_start + j
|
||||
if kv_pos < copy_kv_len:
|
||||
kv_shared[j, d] = v[
|
||||
kv_start + kv_pos, kv_head, d
|
||||
]
|
||||
else:
|
||||
kv_shared[j, d] = T.cast(0, T.bfloat16)
|
||||
|
||||
# 累积当前 tile 的输出贡献：output_accum += probs @ V。
|
||||
T.gemm(
|
||||
probs,
|
||||
kv_shared,
|
||||
output_accum,
|
||||
policy=T.GemmWarpPolicy.FullRow,
|
||||
)
|
||||
|
||||
if valid_q_tile:
|
||||
# 所有 KV tiles 处理完后再归一化。row_scale 已结束循环内使命，
|
||||
# 现在复用为每行的 1/softmax_denominator。每行只除一次，后续
|
||||
# 128 个输出维度都通过乘法归一化。补齐行和完全 mask 行设为0。
|
||||
for i in T.Parallel(packed_block_m):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
row_scale[i] = T.if_then_else(
|
||||
(packed_row < packed_q_len)
|
||||
and (
|
||||
(not causal)
|
||||
or q_pos + causal_offset >= 0
|
||||
),
|
||||
1.0 / row_denom[i],
|
||||
0.0,
|
||||
)
|
||||
|
||||
# 将 packed row 重新映射到原始 output[q_pos, qo_head, d]。
|
||||
# packed tail 只参与内部补齐，不允许写回输出张量。
|
||||
for i, d in T.Parallel(packed_block_m, head_dim_vo):
|
||||
packed_row = packed_tile_start + i
|
||||
q_pos = packed_row // group_size
|
||||
group_head = packed_row % group_size
|
||||
if packed_row < packed_q_len:
|
||||
output[
|
||||
q_start + q_pos,
|
||||
kv_head * group_size + group_head,
|
||||
d,
|
||||
] = output_accum[i, d] * row_scale[i]
|
||||
|
||||
return packed_kernel_factor_softmax_scale_v18
|
||||
|
||||
|
||||
def run_kernel(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
output,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
):
|
||||
"""OJ 调用入口：按完整规格选择、缓存并启动 TileLang kernel。"""
|
||||
|
||||
# cache key 必须包含所有会改变生成代码、网格或张量 shape 的参数：
|
||||
# q/k 的总长度决定静态张量形状；seq_len 决定网格上界、tile 配置和分支；
|
||||
# head 数、head_dim、causal 则直接影响 GQA 映射和 attention 数学逻辑。
|
||||
key = (
|
||||
"packed-gqa-factor-softmax-scale-v18",
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
if key not in _kernel_cache:
|
||||
# 首次遇到该参数组合时才执行 JIT：
|
||||
# seq_len<=128：小序列回退路径，每 CTA 处理一个 Q head，启动和资源
|
||||
# 开销更小，保留单 token、非2次幂尾部的现有优势；
|
||||
# seq_len>128 ：主 packed GQA 路径，在一个 CTA 内处理同一 KV head
|
||||
# 对应的多个 Q heads，提高 K/V 数据复用。
|
||||
if seq_len <= 128:
|
||||
_kernel_cache[key] = build_kernel(
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
else:
|
||||
_kernel_cache[key] = build_packed_kernel_factor_softmax_scale_v18(
|
||||
q.shape[0],
|
||||
k.shape[0],
|
||||
batch_size,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim_qk,
|
||||
head_dim_vo,
|
||||
causal,
|
||||
)
|
||||
# JIT 结果是可直接接收 Torch tensors 的 Cython backend callable。
|
||||
# output 由调用方提前分配，本函数只启动 kernel，不在计时路径创建临时张量。
|
||||
_kernel_cache[key](q, k, v, output, qo_indptr, kv_indptr)
|
||||
|
|
@ -43,6 +43,37 @@ The JIT cache key is:
|
|||
Different indptr distributions with the same compile-time specification reuse
|
||||
the same compiled kernel. Per-request lengths and masks remain runtime values.
|
||||
|
||||
## Optimization version summary
|
||||
|
||||
| Version | Optimization direction | Result and decision |
|
||||
| --- | --- | --- |
|
||||
| opt_000 | FlashAttention-style tiled structural baseline using the default TVM-FFI backend | Representative local correctness passed, but OJ failed during host-library export; retained only as the structural/failure baseline |
|
||||
| opt_001 | Switch only the execution adapter to `execution_backend="cython"` | First OJ-compatible version; 15/15 passed, 690.101 ms; retained as the valid compatibility baseline |
|
||||
| opt_002 | Runtime invalid-Q-tile guard and zero-trip KV loop | Large ragged/Q-shorter gains, but dense cases regressed 10%-12%; retained as evidence and superseded by compile-time specialization |
|
||||
| opt_003 | Compile-time dense/ragged guard, long-path N=32, remove `O_shared`, enable fast math | 15/15 passed, 300.980 ms; major resource/occupancy improvement, retained |
|
||||
| opt_004 | C500-style packed GQA ownership, M128/N64, 512 threads, full-visible causal fast path | Generated 69,632-byte dynamic shared allocation exceeded the 64 KiB limit; rejected before execution |
|
||||
| opt_005 | Reduce packed KV tile N64->N32 to fit shared memory | OJ reused stale opt_004 code, so the intended kernel was not evaluated; inconclusive |
|
||||
| opt_006 | Change builder, PrimFunc, cache key, and local tuning identity to force fresh 128x32 compilation | 15/15 passed, 100.631 ms; packed GQA/KV reuse produced a 2.99x gain over opt_003, retained |
|
||||
| opt_007 | Restore N=64 by lifetime-sharing one shared allocation between K and V | 15/15 passed, 89.595 ms; reduced KV/softmax loop count and improved opt_006 by 10.97%, retained |
|
||||
| opt_008 | Use M64/N32 and 256 threads globally to target two resident CTAs | Resources allowed two CTAs, but doubled KV/softmax iterations; local total regressed 5.57%, rejected |
|
||||
| opt_009 | Shape dispatch: resident-two M64/N32 only for equal lengths up to 1024, otherwise opt_007 M128/N64 | 15/15 OJ passed, 89.365 ms; preserved long performance and improved selected short/equal shapes, retained |
|
||||
| opt_010 | Explicit shared-memory swizzle annotations on the long path | Generated host/device kernels were identical to opt_009; measured as a no-op, rejected |
|
||||
| opt_011 | Reverse dense causal packed-Q tile order so heavy late-Q CTAs are submitted first | All targeted dense cases improved; local total 88.646->87.310 ms (-1.51%), retained |
|
||||
| opt_012 | Reuse softmax scratch, compute one reciprocal per row, multiply during store, and specialize dense equal-length bounds | 15/15 OJ passed; corrected CSV total 84.311 ms and mean score ratio 40.633%; current best |
|
||||
| opt_013 | Abandoned experimental dispatch candidate | Explicitly discarded; its source/results must not be used as optimization evidence or as a parent version |
|
||||
| opt_014 | Change only long-path PV GEMM from FullRow to FullCol to improve V shared reads | Rejected during case-3 compilation: QK scores/probs use a FullRow fragment layout incompatible with FullCol PV |
|
||||
| opt_015 | Move V copy immediately after QK max reduction | 15/15 local correctness passed, but generated case-4 loop gained a sixth barrier and regressed 1.95%; rejected |
|
||||
| opt_016 | Use Square warp partition for both long-path QK and PV | Rejected during case-4 compilation: QK N=64 and PV N=128 still infer incompatible scores/probs layouts |
|
||||
| opt_017 | Skip output rescaling when row_scale is exactly 1 | Case-4 output was identical, but branch overhead regressed latency by 0.53%; rejected |
|
||||
| opt_018 | Factor the common softmax scale as `(x-max)*scale` | Case-4 correctness passed, but latency regressed 1.18%; generated code shows no loop instruction-count reduction; rejected |
|
||||
|
||||
The current submission entry `tilelang/run_kernel.py` is byte-identical to
|
||||
`tilelang/opt_012_dense_softmax_cleanup.py`. Future versions start from opt_012.
|
||||
|
||||
The reusable analysis dimensions and profiler decision tree for future rounds
|
||||
are recorded in
|
||||
`workflows/references/tilelang_kernel_optimization_directions.md`.
|
||||
|
||||
## opt_000: structural baseline
|
||||
|
||||
File: `tilelang/opt_000_structural_baseline.py`
|
||||
|
|
@ -912,3 +943,181 @@ the logical Q tile, so the gain is isolated to causal load scheduling.
|
|||
|
||||
Decision: retain opt_011 as the next OJ candidate. Detailed data is in
|
||||
`results/tilelang_16g/opt_011_results_16g.md`.
|
||||
|
||||
## opt_012: dense softmax cleanup
|
||||
|
||||
File: `tilelang/opt_012_dense_softmax_cleanup.py`
|
||||
|
||||
Parent: `opt_011_reverse_dense_tiles.py`
|
||||
|
||||
This version refines the online softmax computation in the dense path,
|
||||
focusing on reducing numerical overhead and improving instruction efficiency.
|
||||
|
||||
Single conceptual change:
|
||||
|
||||
- Cleanup of dense softmax computation path with optimized exp2 usage
|
||||
- Improved numerical stability and reduced computational overhead
|
||||
|
||||
Local correctness: All 15 OJ test cases pass with 100% match rate locally.
|
||||
|
||||
### opt_012 OJ results
|
||||
|
||||
Submitted: 2026-07-16
|
||||
|
||||
All 15 testcases passed correctness.
|
||||
|
||||
| Case | Config | User (ms) | Baseline (ms) | Speedup | Score % |
|
||||
|------|--------|-----------|---------------|---------|---------|
|
||||
| 1 | ragged_b33_16294 | 2.585 | 1.575 | 0.609x | 35.21% |
|
||||
| 2 | equal_b1_s1024 | 0.291 | 0.260 | 0.893x | 46.77% |
|
||||
| 3 | equal_b1_s4096 | 3.015 | 1.655 | 0.549x | 30.71% |
|
||||
| 4 | equal_b1_s16384 | 45.911 | 22.426 | 0.488x | 26.52% |
|
||||
| 5 | equal_b4_s1024 | 0.904 | 0.635 | 0.702x | 39.26% |
|
||||
| 6 | equal_b4_s4096 | 11.688 | 6.064 | 0.519x | 28.65% |
|
||||
| 7 | equal_b16_s1024 | 3.335 | 2.024 | 0.607x | 34.44% |
|
||||
| 8 | equal_b16_s2048 | 12.394 | 6.638 | 0.536x | 30.08% |
|
||||
| 9 | q512_k1024_b4 | 0.715 | 0.538 | 0.752x | 41.47% |
|
||||
| 10 | mixed_b4 | 0.577 | 0.412 | 0.714x | 39.95% |
|
||||
| 11 | q_lt_kv_b2 | 0.701 | 0.410 | 0.585x | 34.39% |
|
||||
| 12 | ragged_b27_12251 | 1.816 | 1.169 | 0.644x | 36.93% |
|
||||
| 13 | short_ragged_969 | 0.308 | 0.151 | 0.490x | 32.12% |
|
||||
| 14 | single_token | 0.018 | 0.106 | 5.889x | 85.49% |
|
||||
| 15 | tail_non_power2 | 0.053 | 0.109 | 2.057x | 67.50% |
|
||||
|
||||
Aggregate performance:
|
||||
- Sum of user times: 84.311 ms
|
||||
- Sum of baseline times: 44.172 ms
|
||||
- Mean score ratio: 40.633%
|
||||
- Display scores: 35, 46, 30, 26, 39, 28, 34, 30, 41, 40, 34, 37, 32, 85, 67
|
||||
- Mean display score: 40.3 / 100
|
||||
|
||||
The three aggregate values above were recomputed directly from all 15 rows in
|
||||
`results/tilelang/tilelang_oj_results_opt_012.csv`; the earlier manually entered
|
||||
totals were incorrect.
|
||||
|
||||
Observations:
|
||||
- All correctness tests passed
|
||||
- Cases 14 and 15 (edge cases) significantly outperform baseline
|
||||
- Long dense sequences (cases 3, 4, 6, 8) are slowest relative to baseline
|
||||
- Overall performance is below baseline on most compute-heavy cases
|
||||
|
||||
Decision: retain as current best TileLang version. Further optimization should
|
||||
focus on improving the long dense sequence path where the performance gap is
|
||||
largest. The baseline FlashInfer implementation uses more sophisticated
|
||||
techniques (partition-KV, async prefetch, architecture-specific MMA) that
|
||||
TileLang does not yet fully replicate.
|
||||
|
||||
## opt_014 to opt_018: profiler-driven case-4 experiments
|
||||
|
||||
Parent for every candidate: `opt_012_dense_softmax_cleanup.py`. None of these
|
||||
experiments replaced `tilelang/run_kernel.py`; opt_012 remains the submission
|
||||
entry and current verified best.
|
||||
|
||||
### Profiler evidence used
|
||||
|
||||
The opt_012 case-4 profile reports 52,042.51 Kcycles, 99.29% AP busy duty,
|
||||
20.14% MMA duty, 74.04% shared-memory efficiency, and 53,248 bytes dynamic
|
||||
shared memory. Compared with opt_007, `wsm_stall` fell 8.20%, but
|
||||
`vls_pipeline_stall` increased by 569.55%. Generated case-4 code shows five
|
||||
`__syncthreads()` calls per KV-loop iteration. QK loads its shared operand with
|
||||
`uint2`, while PV loads V into its B fragment with scalar BF16 accesses.
|
||||
|
||||
Detailed profiler analysis is stored in
|
||||
`results/tilelang_64g/opt_012_case4_mcprofiler_analysis_64g.md`.
|
||||
|
||||
### opt_014: PV FullCol
|
||||
|
||||
File: `tilelang/opt_014_pv_fullcol.py`
|
||||
|
||||
Hypothesis: distributing PV waves along the 128 output columns could make V
|
||||
shared-memory reads more contiguous. The short resident-two path retained
|
||||
FullRow; only the long path used FullCol for PV.
|
||||
|
||||
Cases 1 and 2 passed because they retained FullRow. Case 3 failed during
|
||||
TileLang layout inference before device compilation. QK produces a FullRow
|
||||
`scores/probs` fragment, while FullCol PV requires a different thread layout.
|
||||
Direct fragment copy cannot perform that cross-thread redistribution. A shared
|
||||
staging buffer could transform it, but would exceed the current 53-KiB resource
|
||||
budget. Decision: reject.
|
||||
|
||||
### opt_015: early V copy
|
||||
|
||||
File: `tilelang/opt_015_v_prefetch_after_qk.py`
|
||||
|
||||
Hypothesis: move V copy from after softmax to immediately after QK max reduction,
|
||||
placing the K/V lifetime transition near an existing synchronization and making
|
||||
V ready earlier.
|
||||
|
||||
All 15 local correctness cases passed. Generated case-4 code disproved the
|
||||
barrier-merging hypothesis: the compiler inserted synchronization before V copy
|
||||
and another before the sum reduction, increasing the KV-loop barrier count from
|
||||
five to six. A seven-sample local A/B measured:
|
||||
|
||||
```text
|
||||
opt_012 median: 45.976318 ms
|
||||
opt_015 median: 46.874752 ms
|
||||
regression: 1.95%
|
||||
```
|
||||
|
||||
Decision: reject. Moving a shared overwrite between the two reductions worsens
|
||||
the compiler's dependency schedule.
|
||||
|
||||
### opt_016: Square QK/PV partition
|
||||
|
||||
File: `tilelang/opt_016_square_warp_partition.py`
|
||||
|
||||
Hypothesis: use Square for both QK and PV so their fragment policy remains
|
||||
consistent while distributing waves across M and N.
|
||||
|
||||
Case-4 compilation still failed layout inference. Policy equality is not enough:
|
||||
QK has output shape M128/N64 while PV has M128/N128, so Square computes different
|
||||
wave partitions and incompatible `scores/probs` layouts. Decision: reject.
|
||||
|
||||
### opt_017: conditional output rescale
|
||||
|
||||
File: `tilelang/opt_017_skip_identity_rescale.py`
|
||||
|
||||
When a KV tile does not raise a row's running maximum, `row_scale=exp2(0)=1`.
|
||||
The candidate branched around the 128 output-accumulator multiplications for
|
||||
such rows. Case-4 output was byte-identical to opt_012, but local A/B measured:
|
||||
|
||||
```text
|
||||
opt_012 median: 45.963520 ms
|
||||
opt_017 median: 46.207359 ms
|
||||
regression: 0.53%
|
||||
```
|
||||
|
||||
The per-row predicate/divergence costs more than the saved identity multiplies
|
||||
on this mapping. Decision: reject.
|
||||
|
||||
### opt_018: factor softmax scale
|
||||
|
||||
File: `tilelang/opt_018_factor_softmax_scale.py`
|
||||
|
||||
The packed path rewrites `x*scale-max*scale` as `(x-max)*scale`. Case-4 output
|
||||
matches opt_012 at 100%, with maximum absolute difference 0.00390625 and worst
|
||||
tolerance ratio 0.13355. Ten-sample local A/B measured:
|
||||
|
||||
```text
|
||||
opt_012 median: 46.044737 ms
|
||||
opt_018 median: 46.588417 ms
|
||||
regression: 1.18%
|
||||
```
|
||||
|
||||
Both generated kernels keep the same `(1024,4,1)` grid, 512 threads, 53,248-byte
|
||||
dynamic shared allocation, and five barriers. In opt_012 the compiler can reuse
|
||||
`row_max*scale` outside the 16-score unrolled loop, so factorization does not
|
||||
reduce the per-score instruction count. The new subtract-then-multiply order
|
||||
instead lengthens the dependency chain feeding `exp2f`. Decision: reject.
|
||||
|
||||
Generated artifacts:
|
||||
|
||||
```text
|
||||
results/tilelang_64g/opt_015_case4_device_kernel.cu
|
||||
results/tilelang_64g/opt_015_case4_host_kernel.cu
|
||||
results/tilelang_64g/opt_018_case4_device_kernel.cu
|
||||
results/tilelang_64g/opt_018_case4_host_kernel.cu
|
||||
```
|
||||
|
||||
Final decision for this round: retain opt_012 unchanged. No further optimization
|
||||
candidate was created after opt_018.
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ and hypotheses that still require experiments.
|
|||
| `mxmaca_matrix_instructions.md` | Known and missing matrix-compute information |
|
||||
| `mxcc_compiler_and_profiler.md` | Architecture targets, resource reports, profiling/debug tools |
|
||||
| `mcflashinfer_kernel_notes.md` | MetaX open-source kernel sources to inspect |
|
||||
| `tilelang_kernel_optimization_directions.md` | Roofline、访存生命周期、流水线、MMA、Softmax 和 profiler 指标到代码动作的完整分析框架 |
|
||||
|
||||
Primary sources:
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,264 @@
|
|||
# TileLang 算子优化方向与性能分析框架
|
||||
|
||||
## 1. 目的
|
||||
|
||||
本文用于指导 FlashInfer Ragged Prefill 在 MetaX C500 上的后续 TileLang
|
||||
优化。先建立可量化的性能模型，再使用 mcTracer 和 mcProfiler 定位具体瓶颈，
|
||||
避免仅根据源码直觉同时修改多个变量。
|
||||
|
||||
当前有效基线为 `tilelang/opt_012_dense_softmax_cleanup.py`，公开入口
|
||||
`tilelang/run_kernel.py` 与其完全一致。opt_013 已放弃，不得作为父版本或性能
|
||||
证据。
|
||||
|
||||
## 2. 固定正确性约束
|
||||
|
||||
- Q/K/V/output 为 BF16，主要累加为 FP32。
|
||||
- Ragged NHD layout，32 个 Q heads、4 个 KV heads，GQA group size 为 8。
|
||||
- QK 和 VO head dimension 均为 128。
|
||||
- Bottom-right causal mask。
|
||||
- 必须支持 Q/KV 不等长、非 2 次幂尾部和无效网格 tile。
|
||||
- `run_kernel` 计时路径中不得同步、分配临时张量或重复 JIT。
|
||||
|
||||
任何优化先满足完整正确性，再讨论性能。
|
||||
|
||||
## 3. 算法和任务划分
|
||||
|
||||
优先确认高层映射，而不是先做局部指令调整：
|
||||
|
||||
- 保持 FlashAttention 风格在线 Softmax，禁止物化完整注意力矩阵。
|
||||
- 同一 CTA 内尽量让 8 个 Q heads 共享对应 KV head 的 K/V 数据。
|
||||
- 跳过完全不可见的 causal tiles 和无效 ragged Q tiles。
|
||||
- 根据 tiny、short equal、dense long、ragged、Q<KV 选择编译期特化。
|
||||
- 因果计算量随 Q 位置增加时，优先调度重 tile，减少 kernel 尾部。
|
||||
- 对超长 KV 评估 split-KV；必须同时计入 merge kernel 和额外 workspace。
|
||||
- 评估 persistent CTA 或精确 work planning 是否能减少矩形网格浪费。
|
||||
|
||||
opt_006 的 packed GQA 和 opt_011 的反向 dense tile 调度说明：任务所有权和
|
||||
调度通常比单条算术指令调整具有更高收益。
|
||||
|
||||
## 4. 计算强度和访存强度
|
||||
|
||||
对单个 `M x N` attention tile，`Dq=Dv=128`：
|
||||
|
||||
```text
|
||||
QK FLOPs = 2 * M * N * 128
|
||||
PV FLOPs = 2 * M * N * 128
|
||||
总 FLOPs = 512 * M * N
|
||||
```
|
||||
|
||||
BF16 主输入的最低数据量近似为：
|
||||
|
||||
```text
|
||||
Q = M * 128 * 2 bytes，在整个 KV 循环中复用
|
||||
K = N * 128 * 2 bytes，每个 KV tile 加载
|
||||
V = N * 128 * 2 bytes，每个 KV tile 加载
|
||||
```
|
||||
|
||||
每个代表性 shape 都要计算：
|
||||
|
||||
1. 理论 FLOPs。
|
||||
2. 理论最低 HBM 字节数。
|
||||
3. `FLOPs / bytes` 算术强度。
|
||||
4. 计算 Roofline 与带宽 Roofline。
|
||||
5. profiler 实际 global/L2 流量与理论最低流量的比值。
|
||||
|
||||
长序列理论上计算强度高，但如果 shared stall 或同步使 MMA duty 很低，仍然
|
||||
无法接近计算 Roofline。
|
||||
|
||||
## 5. 变量访问次数与生命周期
|
||||
|
||||
| 数据 | 期望存储层级 | 需要检查的冗余 |
|
||||
| --- | --- | --- |
|
||||
| Q | Global -> Shared -> MMA fragment | 是否每个 KV tile 重复从 shared 读取；部分 Q 能否常驻寄存器 |
|
||||
| K | Global -> Shared -> MMA fragment | 是否在 Q heads、Q tiles 或 CTA 之间重复加载 |
|
||||
| V | Global -> Shared -> MMA fragment | 是否与 K 安全复用 shared allocation；是否重复加载 |
|
||||
| scores | FP32 fragment | 是否重复 clear/fill/copy，是否形成 private spill |
|
||||
| probs | BF16 fragment | FP32->BF16 转换和 fragment layout 转换是否冗余 |
|
||||
| output accumulator | FP32 fragment | 是否占用过多寄存器，是否被重复缩放或写回 |
|
||||
| row max/sum/denom/scale | FP32 fragment | 是否可以复用 scratch，减少 copy/fill/reduce |
|
||||
| indptr 和长度 | 标量/global | Dense shape 是否可在编译期常量化 |
|
||||
|
||||
应为每个 KV-loop iteration 列出 Q/K/V、scores、softmax state 和 output state
|
||||
的读写次数，并与生成代码及 profiler 实际流量交叉验证。
|
||||
|
||||
## 6. 存储层级与资源限制
|
||||
|
||||
TileLang 主要存储对象：
|
||||
|
||||
```text
|
||||
T.alloc_fragment -> 通常映射到寄存器或 MMA fragment
|
||||
T.alloc_local -> 可能映射到寄存器，也可能形成 private memory
|
||||
T.alloc_shared -> C500 WSM/shared memory
|
||||
```
|
||||
|
||||
每个候选必须记录：
|
||||
|
||||
- threads/CTA 和 64-thread waves/CTA。
|
||||
- registers/thread、MT-register occupancy。
|
||||
- static/dynamic shared memory。
|
||||
- private memory per thread 和总量。
|
||||
- 可驻留 CTA/AP 和 waves/AP。
|
||||
- pipeline stage、padding、swizzle 后的资源变化。
|
||||
|
||||
private memory 是每线程逻辑地址空间，通常由设备内存支持并经过缓存；它不等于
|
||||
每次都直接访问 HBM，但经常意味着寄存器溢出，仍应尽量避免。
|
||||
|
||||
当前 opt_012 长路径继承的关键资源基线约为：
|
||||
|
||||
```text
|
||||
512 threads/CTA
|
||||
53,248 bytes dynamic shared
|
||||
255 registers/thread（历史 opt_007 测量）
|
||||
1 CTA/AP（shared memory 限制）
|
||||
```
|
||||
|
||||
任何 tile 或流水线修改都必须先证明没有超过 64 KiB shared 限制，也没有产生
|
||||
不可接受的 private spill。
|
||||
|
||||
## 7. Shared memory 布局与同步
|
||||
|
||||
重点分析：
|
||||
|
||||
- Q/K/V 的 MMA 读取是否发生 bank conflict。
|
||||
- reduction workspace 是否产生 WSM conflict。
|
||||
- K 写入、QK 读取、V 覆盖、PV 读取之间需要多少 barrier。
|
||||
- padding/swizzle 是否真正改变生成的物理地址。
|
||||
- shared allocation 生命周期能否进一步复用。
|
||||
|
||||
opt_010 已证明：显式添加 swizzle annotation 后生成 host/device kernel 与
|
||||
opt_009 完全相同。TileLang 源码出现 swizzle 不等于物理布局变化，必须比较生成
|
||||
代码或 profiler 指标。
|
||||
|
||||
## 8. 流水线排布
|
||||
|
||||
当前 KV tile 的主要顺序为：
|
||||
|
||||
```text
|
||||
load K
|
||||
-> QK GEMM
|
||||
-> max reduction
|
||||
-> exp2
|
||||
-> sum reduction
|
||||
-> rescale denominator/output
|
||||
-> load V
|
||||
-> PV GEMM
|
||||
```
|
||||
|
||||
可能的优化方向：
|
||||
|
||||
- 当前 tile 计算时预取下一 tile 的 K 或 V。
|
||||
- K/V 双缓冲，或仅对 K/仅对 V 做异步预取。
|
||||
- 将 global->shared copy 与 Softmax 标量计算重叠。
|
||||
- 减少每轮 `__syncthreads()`。
|
||||
- 缩小某个 shared tile，为第二 pipeline stage 腾出空间。
|
||||
|
||||
当前 K/V 分时复用同一 shared allocation，在线 Softmax 又存在循环依赖，不能
|
||||
直接把 `NUM_STAGES=1` 改成 2。必须先画出 buffer 生命周期并重新计算 shared
|
||||
总量。
|
||||
|
||||
## 9. 连续访问、合并与向量化
|
||||
|
||||
检查 Global->Shared 和 Shared->Fragment：
|
||||
|
||||
- 相邻线程是否访问连续地址。
|
||||
- 地址是否满足 BF16 向量加载对齐。
|
||||
- 完整 tile 是否走向量化 `T.copy`。
|
||||
- 只有最后一个 ragged tile 使用标量 tail 路径。
|
||||
- Q/K/V 是否生成 16B、32B 或更宽事务。
|
||||
- 输出写回是否合并。
|
||||
- packed-row 映射是否破坏 Q 的连续访问。
|
||||
|
||||
必须从生成代码确认 `uint2`、`uint4` 等实际宽加载；不能只根据 TileLang 源码
|
||||
推断已经向量化。
|
||||
|
||||
## 10. MMA 映射与线程组织
|
||||
|
||||
分析：
|
||||
|
||||
- `BLOCK_M/BLOCK_N/K` 是否匹配 C500 MMA 形状。
|
||||
- `T.GemmWarpPolicy.FullRow` 是否是最佳 policy。
|
||||
- threads 是否为 64 的合理倍数。
|
||||
- 每个 CTA 的 wave 数和每 AP 的驻留 wave 数。
|
||||
- QK 和 PV 是否都高效使用 MMA。
|
||||
- fragment layout 转换是否产生额外指令。
|
||||
- MMA duty 低是因为操作数等待，还是 GEMM 映射本身低效。
|
||||
|
||||
AP busy 高只表示 AP 有工作，不代表矩阵单元得到充分利用。
|
||||
|
||||
## 11. Softmax 与 reduction
|
||||
|
||||
每个 KV tile 通常包含 max reduction、exp2、sum reduction、running state 更新和
|
||||
output rescale。继续检查：
|
||||
|
||||
- max/sum 是否通过 shared workspace，是否可用 wave reduction 减少 barrier。
|
||||
- fully-visible tile 是否完全移除逐元素 mask。
|
||||
- `row_scale`、max 和 denominator 是否重复计算或复制。
|
||||
- FP32/BF16 转换次数。
|
||||
- softmax 循环是否可融合。
|
||||
- 是否存在不必要的 `fill`、`copy`、`clear`。
|
||||
|
||||
opt_012 已完成 scratch 复用、每行一次倒数和乘法归一化，下一轮必须基于生成
|
||||
指令确认还剩哪些重复工作。
|
||||
|
||||
## 12. 循环展开与代码体积
|
||||
|
||||
- 固定且较小的 MMA K-loop、向量 copy loop 可评估展开。
|
||||
- 长 KV loop 禁止完全展开。
|
||||
- Softmax 元素循环过度展开可能增加寄存器、private memory 和指令缓存压力。
|
||||
- 每次展开实验必须重新检查资源报告和生成代码体积。
|
||||
|
||||
## 13. mcTracer 与 mcProfiler 的指标映射
|
||||
|
||||
先用 mcTracer：
|
||||
|
||||
- 确认目标 kernel symbol、grid/block 和真实持续时间。
|
||||
- 获取 registers、dynamic shared、private memory 和 occupancy。
|
||||
- 排除 host launch、同步、额外 memcpy、重编译和其他 kernel 干扰。
|
||||
|
||||
再用 mcProfiler：
|
||||
|
||||
- AP busy、MMA/MTE/STE/VLS/L2C duty。
|
||||
- ISU stall 分类。
|
||||
- shared-memory access efficiency 和 bank conflict。
|
||||
- VL1/L2 hit rate。
|
||||
- global read/write bytes。
|
||||
- achieved/dispatched waves。
|
||||
|
||||
根据结果选择方向：
|
||||
|
||||
| 证据 | 优先代码方向 |
|
||||
| --- | --- |
|
||||
| `wsm_stall` 仍占绝对主导、shared efficiency 低 | Shared layout、reduction workspace、barrier、buffer 生命周期 |
|
||||
| `wsm_stall` 下降但 MMA duty 仍低 | GEMM policy、tile/wave 映射、Softmax 标量指令占比 |
|
||||
| MTE duty 或 memory latency 高 | K/V 向量加载、异步预取、双缓冲、L2 reuse |
|
||||
| private memory 非零或 registers 过高 | 缩小 fragment、减少展开、复用 softmax scratch |
|
||||
| 单 CTA 尾部明显、长 KV 工作不均衡 | CTA 顺序、persistent work planning、split-KV |
|
||||
| host launch 占比高 | 仅针对 tiny path 优化入口和 kernel 数量 |
|
||||
|
||||
## 14. opt_012 当前分析优先级
|
||||
|
||||
首先分析正式 case 4，然后用 case 3/6/8 验证同一结论是否可推广：
|
||||
|
||||
1. `wsm_stall` 是否仍接近 opt_007 的约 98%。
|
||||
2. MMA duty 是否已高于 opt_007 的 15%-18%。
|
||||
3. shared efficiency 是否仍接近 79.62%。
|
||||
4. opt_012 的 softmax cleanup 降低了哪些指令和 stall。
|
||||
5. case 4 当前是 shared/barrier、MMA throughput，还是 load pipeline 受限。
|
||||
|
||||
候选优化按证据排序：
|
||||
|
||||
1. Reduction/barrier 和 shared workspace 优化。
|
||||
2. K/V copy 与计算重叠。
|
||||
3. 保持正确资源预算的双缓冲或部分预取。
|
||||
4. MMA tile/policy 和 wave mapping。
|
||||
5. 超长序列 split-KV。
|
||||
|
||||
## 15. 实验规则
|
||||
|
||||
- 每轮只改变一个可解释的概念，保留父版本。
|
||||
- builder、PrimFunc 和 Python cache key 同时升级版本，避免陈旧缓存。
|
||||
- 先正确性，再资源报告，再局部 A/B，再完整 15 点。
|
||||
- 跨 OJ 版本比较原始 `tk_time_ms`，不使用波动的 baseline 归因。
|
||||
- 小于正常噪声的变化必须重复测量。
|
||||
- profiler 采集脚本不得在同一进程加载 FlashInfer 和 TileLang。
|
||||
|
||||
|
|
@ -0,0 +1,43 @@
|
|||
# DeepSeek MLA Decode TileLang 算子优化
|
||||
|
||||
本目录用于实现和优化 DeepSeek V3/R1 decode 阶段的 Multi-Head Latent
|
||||
Attention（MLA）变体，目标硬件为 MetaX C500，提交实现限定使用 TileLang。
|
||||
|
||||
## 入口
|
||||
|
||||
- `problem_statement.md`：整理后的题目契约、数学定义、接口和测试范围。
|
||||
- `赛题一mla教程.md`：已有提交说明、当前约 49.5 分实现和完整示例代码。
|
||||
- `test_tilelang_mla.py`：已有 TileLang kernel、PyTorch reference 和本地测试入口。
|
||||
- `test_cases_mla_batch_ctx.json`：31 个 batch/kv_ctx 测试规格。
|
||||
|
||||
## 当前状态
|
||||
|
||||
已有实现采用：
|
||||
|
||||
```text
|
||||
BLOCK_H = 16
|
||||
BLOCK_N = 32
|
||||
num_split = 1
|
||||
FP16 input/output
|
||||
FP32 QK、online softmax 和 output accumulation
|
||||
FullCol GEMM policy
|
||||
```
|
||||
|
||||
已有教程记录一次 Accepted、约 49.5 分结果。该结果作为后续优化的初始有效
|
||||
基线；开始修改前应重新确认提交源码、OJ 原始数据和当前容器中的生成代码一致。
|
||||
|
||||
## 工作规则
|
||||
|
||||
- `run_kernel` 的函数名、参数顺序和参数类型不得改变。
|
||||
- 提交 kernel 内不得使用 PyTorch 完成 GPU 计算。
|
||||
- 结果必须写入调用方提供的 `output`。
|
||||
- 当前最佳版本必须独立保留；每个候选使用新文件和唯一 JIT cache identity。
|
||||
- 正确性优先，容差为 `rtol=2e-3, atol=2e-3`。
|
||||
- 长上下文可能接近 4 GiB 评测内存限制，workspace 必须纳入预算。
|
||||
|
||||
通用优化方法参考：
|
||||
|
||||
```text
|
||||
/data/operator_task_package/ref/README.md
|
||||
/data/operator_task_package/ref/08_templates/new_operator_checklist.md
|
||||
```
|
||||
Binary file not shown.
|
|
@ -0,0 +1,174 @@
|
|||
# TileLang 算子优化：DeepSeek MLA Decode
|
||||
|
||||
## 1. 任务
|
||||
|
||||
实现 DeepSeek V3/R1 decode 阶段使用的 Multi-Head Latent Attention 变体。
|
||||
评测程序在 GPU 上构造输入，调用提交文件中的 `run_kernel`，实现必须使用
|
||||
TileLang 完成计算并将结果写入 `output`。
|
||||
|
||||
时间限制：10,000 ms。内存限制：4,096 MiB。
|
||||
|
||||
## 2. 固定参数
|
||||
|
||||
| 参数 | 值 |
|
||||
|---|---:|
|
||||
| query/key 非 RoPE 维度 `dim` | 512 |
|
||||
| RoPE 维度 `pe_dim` | 64 |
|
||||
| QK 总维度 | 576 |
|
||||
| V/output 维度 | 512 |
|
||||
| query heads | 16（当前评测范围） |
|
||||
| KV heads | 1 |
|
||||
| softmax scale | `1 / sqrt(576)` |
|
||||
| 输入/输出 dtype | float16 |
|
||||
| 主要数值累加 | 建议 float32 |
|
||||
|
||||
所有 query heads 共享同一组 latent KV cache。
|
||||
|
||||
## 3. 数学定义
|
||||
|
||||
对 batch `b`、query head `h`、KV 位置 `s`：
|
||||
|
||||
```text
|
||||
score[b,h,s] =
|
||||
sum(d=0..511) q[b,h,d] * kv[b,s,0,d]
|
||||
+ sum(p=0..63) q_pe[b,h,p] * k_pe[b,s,0,p]
|
||||
|
||||
scaled_score[b,h,s] = score[b,h,s] / sqrt(576)
|
||||
|
||||
attention[b,h,:] = softmax(scaled_score[b,h,:], dim=kv_ctx)
|
||||
|
||||
output[b,h,d] =
|
||||
sum(s=0..kv_ctx-1) attention[b,h,s] * kv[b,s,0,d]
|
||||
```
|
||||
|
||||
等价矩阵形式：
|
||||
|
||||
```text
|
||||
Q_full = concat(q, q_pe) # (..., 576)
|
||||
K_full = concat(kv, k_pe) # (..., kv_ctx, 576)
|
||||
S = Q_full @ K_full^T / sqrt(576)
|
||||
P = softmax(S, axis=-1)
|
||||
O = P @ kv # (..., 512)
|
||||
```
|
||||
|
||||
不得物化完整大规模 score/attention 到 PyTorch tensor；TileLang kernel 可使用
|
||||
tiled online softmax 或 split/merge 实现。
|
||||
|
||||
## 4. 接口契约
|
||||
|
||||
提交 Python 文件必须提供以下函数，名称和参数顺序完全一致：
|
||||
|
||||
```python
|
||||
def run_kernel(
|
||||
q, # Tensor[float16], (batch, heads, dim)
|
||||
q_pe, # Tensor[float16], (batch, heads, pe_dim)
|
||||
kv, # Tensor[float16], (batch, kv_ctx, kv_heads, dim)
|
||||
k_pe, # Tensor[float16], (batch, kv_ctx, kv_heads, pe_dim)
|
||||
output, # Tensor[float16], (batch, heads, dim)
|
||||
batch, # int64
|
||||
heads, # int64
|
||||
kv_heads, # int64
|
||||
kv_ctx, # int64
|
||||
dim, # int64, fixed 512
|
||||
pe_dim, # int64, fixed 64
|
||||
):
|
||||
...
|
||||
```
|
||||
|
||||
张量均为连续 float16 CUDA tensor。`run_kernel` 自行选择 launch 配置、取得或
|
||||
编译 TileLang kernel，并写入已分配的 `output`。
|
||||
|
||||
## 5. 正确性
|
||||
|
||||
评测使用：
|
||||
|
||||
```text
|
||||
torch.allclose(candidate, reference, rtol=2e-3, atol=2e-3)
|
||||
```
|
||||
|
||||
该容差比此前 Ragged Prefill 案例更严格。online softmax 的 running max、分母、
|
||||
输出分子以及 QK/PV accumulation 应优先使用 float32，并检查长上下文数值稳定性。
|
||||
|
||||
## 6. 测试范围
|
||||
|
||||
```text
|
||||
batch: 1, 2, 4, 8, 16, 32
|
||||
heads: 16
|
||||
kv_heads: 1
|
||||
kv_ctx: 2,048 到 65,536
|
||||
dim: 512
|
||||
pe_dim: 64
|
||||
```
|
||||
|
||||
本目录的 `test_cases_mla_batch_ctx.json` 当前包含 31 个组合：
|
||||
|
||||
| Batch | kv_ctx |
|
||||
|---:|---|
|
||||
| 1 | 2048, 4096, 8192, 16384, 32768, 65536 |
|
||||
| 2 | 2048, 8192, 16384, 32768, 65536 |
|
||||
| 4 | 2048, 8192, 16384, 32768, 65536 |
|
||||
| 8 | 2048, 8192, 16384, 32768, 65536 |
|
||||
| 16 | 2048, 8192, 16384, 32768, 65536 |
|
||||
| 32 | 2048, 8192, 16384, 32768, 65536 |
|
||||
|
||||
大上下文显存占用高，评测已减少 warmup 和 iteration 数量。不能仅依据小 case
|
||||
选择 split 数、tile 或 workspace。
|
||||
|
||||
## 7. 计算与数据规模
|
||||
|
||||
忽略 softmax 标量操作，近似 FLOPs：
|
||||
|
||||
```text
|
||||
QK main = 2 * batch * heads * kv_ctx * 512
|
||||
QK PE = 2 * batch * heads * kv_ctx * 64
|
||||
PV = 2 * batch * heads * kv_ctx * 512
|
||||
Total = 2 * batch * heads * kv_ctx * 1088
|
||||
```
|
||||
|
||||
当 `heads=16`：
|
||||
|
||||
```text
|
||||
Total FLOPs = 34,816 * batch * kv_ctx
|
||||
```
|
||||
|
||||
每个 batch 的共享 cache 最低输入字节近似：
|
||||
|
||||
```text
|
||||
KV + K_PE = kv_ctx * (512 + 64) * 2
|
||||
= kv_ctx * 1,152 bytes
|
||||
```
|
||||
|
||||
由于 16 个 heads 共享 KV，CTA ownership 是否能让多 heads 复用 KV/K_PE 是核心
|
||||
性能变量。若每个 head 独立扫描，cache 数据会被重复读取；若一次处理多个 heads，
|
||||
则会增加 Q/output fragment 和 shared/register 压力。
|
||||
|
||||
最大组合 `batch=32, kv_ctx=65536` 中，仅 `kv` 输入约为 2 GiB，`k_pe` 约为
|
||||
256 MiB；split partial output、LSE 和其他 workspace 必须控制在 4 GiB 总限制内。
|
||||
|
||||
## 8. 禁止事项与提交要求
|
||||
|
||||
- GPU 计算算子不得使用 PyTorch，只允许 TileLang 实现计算。
|
||||
- 不得改变 `run_kernel` 接口。
|
||||
- 不得在计时路径中做不必要同步、重复 JIT 或大规模临时分配。
|
||||
- 不得假设 `kv_ctx` 只取某一个长度。
|
||||
- 提交内容应是评测指南要求的 Python 源码，而不是测试脚本或 reference。
|
||||
|
||||
## 9. PyTorch reference 语义
|
||||
|
||||
题目 reference 的逻辑为：
|
||||
|
||||
```python
|
||||
group_num = heads // kv_heads
|
||||
q_main = q.reshape(batch, kv_heads, group_num, dim).permute(0, 2, 1, 3).float()
|
||||
q_pos = q_pe.reshape(batch, kv_heads, group_num, pe_dim).permute(0, 2, 1, 3).float()
|
||||
kv_main = kv.permute(0, 2, 1, 3).float()
|
||||
k_pos = k_pe.permute(0, 2, 1, 3).float()
|
||||
query = torch.cat([q_main, q_pos], dim=-1)
|
||||
key = torch.cat([kv_main, k_pos], dim=-1)
|
||||
scores = torch.einsum("bghd,bhsd->bghs", query, key)
|
||||
attention = torch.softmax(scores * ((dim + pe_dim) ** -0.5), dim=-1)
|
||||
out = torch.einsum("bghs,bhsd->bghd", attention, kv_main)
|
||||
output.copy_(out.permute(0, 2, 1, 3).reshape(batch, heads, dim).to(output.dtype))
|
||||
```
|
||||
|
||||
这段代码只用于理解语义和本地 reference，不可作为提交 kernel 的计算实现。
|
||||
|
|
@ -0,0 +1,33 @@
|
|||
[
|
||||
{"case_id": 1, "batch": 1, "kv_ctx": 2048, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 2, "batch": 1, "kv_ctx": 4096, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 3, "batch": 1, "kv_ctx": 8192, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 4, "batch": 1, "kv_ctx": 16384, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 5, "batch": 1, "kv_ctx": 32768, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 6, "batch": 1, "kv_ctx": 65536, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 7, "batch": 2, "kv_ctx": 2048, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 8, "batch": 2, "kv_ctx": 8192, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 9, "batch": 2, "kv_ctx": 16384, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 10, "batch": 2, "kv_ctx": 32768, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 11, "batch": 2, "kv_ctx": 65536, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 12, "batch": 4, "kv_ctx": 2048, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 13, "batch": 4, "kv_ctx": 8192, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 14, "batch": 4, "kv_ctx": 16384, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 15, "batch": 4, "kv_ctx": 32768, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 16, "batch": 4, "kv_ctx": 65536, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 17, "batch": 8, "kv_ctx": 2048, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 18, "batch": 8, "kv_ctx": 8192, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 19, "batch": 8, "kv_ctx": 16384, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 20, "batch": 8, "kv_ctx": 32768, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 21, "batch": 8, "kv_ctx": 65536, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 22, "batch": 16, "kv_ctx": 2048, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 23, "batch": 16, "kv_ctx": 8192, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 24, "batch": 16, "kv_ctx": 16384, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 25, "batch": 16, "kv_ctx": 32768, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 26, "batch": 16, "kv_ctx": 65536, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 27, "batch": 32, "kv_ctx": 2048, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 28, "batch": 32, "kv_ctx": 8192, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 29, "batch": 32, "kv_ctx": 16384, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 30, "batch": 32, "kv_ctx": 32768, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64},
|
||||
{"case_id": 31, "batch": 32, "kv_ctx": 65536, "heads": 16, "kv_heads": 1, "dim": 512, "pe_dim": 64}
|
||||
]
|
||||
|
|
@ -0,0 +1,290 @@
|
|||
import torch
|
||||
import torch.nn.functional as F
|
||||
import tilelang
|
||||
from tilelang.autotuner import *
|
||||
import tilelang.language as T
|
||||
from einops import rearrange, einsum
|
||||
import argparse
|
||||
import json
|
||||
@tilelang.jit(
|
||||
out_idx=[4],
|
||||
pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True},
|
||||
)
|
||||
def flashattn(batch, heads, kv_head_num, seqlen_kv, dim, pe_dim, block_N, block_H, num_split, softmax_scale):
|
||||
scale = float(softmax_scale * 1.44269504) # log2(e)
|
||||
dtype = T.float16
|
||||
accum_dtype = T.float32
|
||||
kv_group_num = heads // kv_head_num
|
||||
VALID_BLOCK_H = min(block_H, kv_group_num)
|
||||
assert kv_head_num == 1, "kv_head_num must be 1"
|
||||
@T.prim_func
|
||||
def main_split(
|
||||
Q: T.Tensor([batch, heads, dim], dtype),
|
||||
Q_pe: T.Tensor([batch, heads, pe_dim], dtype),
|
||||
KV: T.Tensor([batch, seqlen_kv, kv_head_num, dim], dtype),
|
||||
K_pe: T.Tensor([batch, seqlen_kv, kv_head_num, pe_dim], dtype),
|
||||
Output: T.Tensor([batch, heads, dim], dtype),
|
||||
):
|
||||
glse = T.alloc_global([batch, heads, num_split], dtype)
|
||||
Output_partial = T.alloc_global([batch, heads, num_split, dim], dtype)
|
||||
# flash_attn_split
|
||||
with T.Kernel(batch, heads // min(block_H, kv_group_num), num_split, threads=256) as (bid, hid, bz):
|
||||
Q_shared = T.alloc_shared([block_H, dim], dtype)
|
||||
S_shared = T.alloc_shared([block_H, block_N], dtype)
|
||||
Q_pe_shared = T.alloc_shared([block_H, pe_dim], dtype)
|
||||
KV_shared = T.alloc_shared([block_N, dim], dtype)
|
||||
K_pe_shared = T.alloc_shared([block_N, pe_dim], dtype)
|
||||
O_shared = T.alloc_shared([block_H, dim], dtype)
|
||||
acc_s = T.alloc_fragment([block_H, block_N], accum_dtype)
|
||||
acc_s_cast = T.alloc_fragment([block_H, block_N], dtype)
|
||||
acc_o = T.alloc_fragment([block_H, dim], accum_dtype)
|
||||
scores_max = T.alloc_fragment([block_H], accum_dtype)
|
||||
scores_max_prev = T.alloc_fragment([block_H], accum_dtype)
|
||||
scores_scale = T.alloc_fragment([block_H], accum_dtype)
|
||||
scores_sum = T.alloc_fragment([block_H], accum_dtype)
|
||||
logsum = T.alloc_fragment([block_H], accum_dtype)
|
||||
cur_kv_head = hid // (kv_group_num // block_H)
|
||||
T.use_swizzle(10)
|
||||
T.copy(Q[bid, hid * VALID_BLOCK_H : (hid + 1) * VALID_BLOCK_H, :], Q_shared)
|
||||
T.copy(Q_pe[bid, hid * VALID_BLOCK_H : (hid + 1) * VALID_BLOCK_H, :], Q_pe_shared)
|
||||
T.fill(acc_o, 0)
|
||||
T.fill(logsum, 0)
|
||||
T.fill(scores_max, -T.infinity(accum_dtype))
|
||||
loop_range = T.ceildiv((seqlen_kv // num_split), block_N)
|
||||
for k in T.Pipelined(loop_range, num_stages=2):
|
||||
kv_start = (seqlen_kv // num_split) * bz + k * block_N
|
||||
kv_end = (seqlen_kv // num_split) * bz + (k + 1) * block_N
|
||||
T.copy(KV[bid, kv_start:kv_end, cur_kv_head, :], KV_shared)
|
||||
T.copy(K_pe[bid, kv_start:kv_end, cur_kv_head, :], K_pe_shared)
|
||||
T.clear(acc_s)
|
||||
T.gemm(Q_shared, KV_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol)
|
||||
T.gemm(Q_pe_shared, K_pe_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol)
|
||||
T.copy(scores_max, scores_max_prev)
|
||||
T.fill(scores_max, -T.infinity(accum_dtype))
|
||||
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
|
||||
for i in T.Parallel(block_H):
|
||||
scores_max[i] = T.max(scores_max[i], scores_max_prev[i])
|
||||
for i in T.Parallel(block_H):
|
||||
scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
|
||||
for i, j in T.Parallel(block_H, block_N):
|
||||
acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
|
||||
T.reduce_sum(acc_s, scores_sum, dim=1)
|
||||
T.copy(acc_s, S_shared)
|
||||
T.copy(S_shared, acc_s_cast)
|
||||
for i in T.Parallel(block_H):
|
||||
logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
|
||||
for i, j in T.Parallel(block_H, dim):
|
||||
acc_o[i, j] *= scores_scale[i]
|
||||
T.gemm(acc_s_cast, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol)
|
||||
for i, j in T.Parallel(block_H, dim):
|
||||
acc_o[i, j] /= logsum[i]
|
||||
for i in T.Parallel(block_H):
|
||||
logsum[i] = T.log2(logsum[i]) + scores_max[i] * scale
|
||||
T.copy(logsum, glse[bid, hid * VALID_BLOCK_H : (hid + 1) * VALID_BLOCK_H, bz])
|
||||
T.copy(acc_o, O_shared)
|
||||
T.copy(O_shared, Output_partial[bid, hid * VALID_BLOCK_H : (hid + 1) * VALID_BLOCK_H, bz, :])
|
||||
# combine
|
||||
with T.Kernel(heads, batch, threads=128) as (hid, bz):
|
||||
po_local = T.alloc_fragment([dim], dtype)
|
||||
o_accum_local = T.alloc_fragment([dim], accum_dtype)
|
||||
lse_local_split = T.alloc_var(accum_dtype)
|
||||
lse_logsum_local = T.alloc_var(accum_dtype)
|
||||
lse_max_local = T.alloc_var(accum_dtype)
|
||||
scale_local = T.alloc_var(accum_dtype)
|
||||
T.clear(lse_logsum_local)
|
||||
T.clear(o_accum_local)
|
||||
lse_max_local = -T.infinity(accum_dtype)
|
||||
for k in T.serial(num_split):
|
||||
lse_max_local = T.max(lse_max_local, glse[bz, hid, k])
|
||||
for k in T.Pipelined(num_split, num_stages=1):
|
||||
lse_local_split = glse[bz, hid, k]
|
||||
lse_logsum_local += T.exp2(lse_local_split - lse_max_local)
|
||||
lse_logsum_local = T.log2(lse_logsum_local) + lse_max_local
|
||||
for k in T.serial(num_split):
|
||||
for i in T.Parallel(dim):
|
||||
po_local[i] = Output_partial[bz, hid, k, i]
|
||||
lse_local_split = glse[bz, hid, k]
|
||||
scale_local = T.exp2(lse_local_split - lse_logsum_local)
|
||||
for i in T.Parallel(dim):
|
||||
o_accum_local[i] += po_local[i] * scale_local
|
||||
for i in T.Parallel(dim):
|
||||
Output[bz, hid, i] = o_accum_local[i]
|
||||
@T.prim_func
|
||||
def main_no_split(
|
||||
Q: T.Tensor([batch, heads, dim], dtype),
|
||||
Q_pe: T.Tensor([batch, heads, pe_dim], dtype),
|
||||
KV: T.Tensor([batch, seqlen_kv, kv_head_num, dim], dtype),
|
||||
K_pe: T.Tensor([batch, seqlen_kv, kv_head_num, pe_dim], dtype),
|
||||
Output: T.Tensor([batch, heads, dim], dtype),
|
||||
):
|
||||
with T.Kernel(heads // min(block_H, kv_group_num), batch, threads=128) as (hid, bid):
|
||||
Q_shared = T.alloc_shared([block_H, dim], dtype)
|
||||
S_shared = T.alloc_shared([block_H, block_N], dtype)
|
||||
Q_pe_shared = T.alloc_shared([block_H, pe_dim], dtype)
|
||||
KV_shared = T.alloc_shared([block_N, dim], dtype)
|
||||
K_pe_shared = T.alloc_shared([block_N, pe_dim], dtype)
|
||||
O_shared = T.alloc_shared([block_H, dim], dtype)
|
||||
acc_s = T.alloc_fragment([block_H, block_N], accum_dtype)
|
||||
acc_o = T.alloc_fragment([block_H, dim], accum_dtype)
|
||||
scores_max = T.alloc_fragment([block_H], accum_dtype)
|
||||
scores_max_prev = T.alloc_fragment([block_H], accum_dtype)
|
||||
scores_scale = T.alloc_fragment([block_H], accum_dtype)
|
||||
scores_sum = T.alloc_fragment([block_H], accum_dtype)
|
||||
logsum = T.alloc_fragment([block_H], accum_dtype)
|
||||
cur_kv_head = hid // (kv_group_num // block_H)
|
||||
T.copy(Q[bid, hid * VALID_BLOCK_H : (hid + 1) * VALID_BLOCK_H, :], Q_shared)
|
||||
T.copy(Q_pe[bid, hid * VALID_BLOCK_H : (hid + 1) * VALID_BLOCK_H, :], Q_pe_shared)
|
||||
T.fill(acc_o, 0)
|
||||
T.fill(logsum, 0)
|
||||
T.fill(scores_max, -T.infinity(accum_dtype))
|
||||
loop_range = T.ceildiv(seqlen_kv, block_N)
|
||||
for k in T.Pipelined(loop_range, num_stages=0):
|
||||
T.copy(KV[bid, k * block_N : (k + 1) * block_N, cur_kv_head, :], KV_shared)
|
||||
T.copy(K_pe[bid, k * block_N : (k + 1) * block_N, cur_kv_head, :], K_pe_shared)
|
||||
T.gemm(Q_shared, KV_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol, clear_accum=True)
|
||||
T.gemm(Q_pe_shared, K_pe_shared, acc_s, transpose_B=True, policy=T.GemmWarpPolicy.FullCol)
|
||||
T.copy(scores_max, scores_max_prev)
|
||||
T.fill(scores_max, -T.infinity(accum_dtype))
|
||||
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
|
||||
for i in T.Parallel(block_H):
|
||||
scores_max[i] = T.max(scores_max[i], scores_max_prev[i])
|
||||
for i in T.Parallel(block_H):
|
||||
scores_scale[i] = T.exp2(scores_max_prev[i] * scale - scores_max[i] * scale)
|
||||
for i, j in T.Parallel(block_H, block_N):
|
||||
acc_s[i, j] = T.exp2(acc_s[i, j] * scale - scores_max[i] * scale)
|
||||
T.reduce_sum(acc_s, scores_sum, dim=1)
|
||||
T.copy(acc_s, S_shared)
|
||||
for i in T.Parallel(block_H):
|
||||
logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
|
||||
for i, j in T.Parallel(block_H, dim):
|
||||
acc_o[i, j] *= scores_scale[i]
|
||||
T.gemm(S_shared, KV_shared, acc_o, policy=T.GemmWarpPolicy.FullCol)
|
||||
for i, j in T.Parallel(block_H, dim):
|
||||
acc_o[i, j] /= logsum[i]
|
||||
T.copy(acc_o, O_shared)
|
||||
T.copy(O_shared, Output[bid, hid * VALID_BLOCK_H : (hid + 1) * VALID_BLOCK_H, :])
|
||||
if num_split > 1:
|
||||
return main_split
|
||||
else:
|
||||
return main_no_split
|
||||
def ref_program(q, q_pe, kv, k_pe):
|
||||
# """
|
||||
# Inputs:
|
||||
# - q (Tensor): [batch, heads, dim]
|
||||
# - q_pe (Tensor): [batch, heads, pe_dim]
|
||||
# - kv (Tensor): [batch, seqlen_kv, kv_head_num, dim]
|
||||
# - k_pe (Tensor): [batch, seqlen_kv, kv_head_num, pe_dim]
|
||||
# Outputs:
|
||||
# - output (Tensor): [batch, heads, dim]
|
||||
# """
|
||||
dim = q.shape[-1]
|
||||
pe_dim = q_pe.shape[-1]
|
||||
num_head_groups = q.shape[1] // kv.shape[2]
|
||||
scale = (dim + pe_dim) ** 0.5
|
||||
q = rearrange(q, "b (h g) d -> b g h d", g=num_head_groups) # [batch_size, num_head_groups, groups, dim]
|
||||
q_pe = rearrange(q_pe, "b (h g) d -> b g h d", g=num_head_groups) # [batch_size, num_head_groups, groups, pe_dim]
|
||||
kv = rearrange(kv, "b n h d -> b h n d") # [batch_size, groups, seqlen_kv, dim]
|
||||
k_pe = rearrange(k_pe, "b n h d -> b h n d") # [batch_size, num_head_groups, groups, pe_dim]
|
||||
query = torch.concat([q, q_pe], dim=-1)
|
||||
key = torch.concat([kv, k_pe], dim=-1)
|
||||
scores = einsum(query, key, "b g h d, b h s d -> b g h s") # [batch_size, num_head_groups, groups, seqlen_kv]
|
||||
attention = F.softmax(scores / scale, dim=-1) # [batch_size, num_head_groups, groups, seqlen_kv]
|
||||
out = einsum(attention, kv, "b g h s, b h s d -> b g h d") # [batch_size, num_head_groups, groups, dim]
|
||||
out = rearrange(out, "b g h d -> b (h g) d") # [batch_size, heads, dim]
|
||||
return out
|
||||
def main(
|
||||
batch=1,
|
||||
heads=128,
|
||||
kv_heads=1,
|
||||
kv_ctx=8192,
|
||||
dim=512,
|
||||
pe_dim=64,
|
||||
):
|
||||
qk_flops = 2 * batch * heads * kv_ctx * (dim + pe_dim)
|
||||
pv_flops = 2 * batch * heads * kv_ctx * dim
|
||||
total_flops = qk_flops + pv_flops
|
||||
BLOCK_N = 32
|
||||
BLOCK_H = min(16, heads // kv_heads)
|
||||
num_split = 1
|
||||
softmax_scale = (dim + pe_dim) ** -0.5
|
||||
kernel = flashattn(batch, heads, kv_heads, kv_ctx, dim, pe_dim, BLOCK_N, BLOCK_H, num_split, softmax_scale)
|
||||
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Randn)
|
||||
profiler.assert_allclose(ref_program, rtol=2e-4, atol=1e-4)
|
||||
latency = profiler.do_bench(warmup=500)
|
||||
tflops = total_flops / latency * 1e-9
|
||||
print(f"Latency: {latency} ms")
|
||||
print(f"TFlops: {tflops} TFlops")
|
||||
return latency, tflops
|
||||
def run_regression_perf(
|
||||
batch=1,
|
||||
heads=128,
|
||||
kv_heads=1,
|
||||
kv_ctx=8192,
|
||||
dim=512,
|
||||
pe_dim=64,
|
||||
):
|
||||
BLOCK_N = 64
|
||||
BLOCK_H = min(64, heads // kv_heads)
|
||||
num_split = 1
|
||||
softmax_scale = (dim + pe_dim) ** -0.5
|
||||
kernel = flashattn(batch, heads, kv_heads, kv_ctx, dim, pe_dim, BLOCK_N, BLOCK_H, num_split, softmax_scale)
|
||||
profiler = kernel.get_profiler(tensor_supply_type=tilelang.TensorSupplyType.Randn)
|
||||
profiler.assert_allclose(ref_program, rtol=2e-4, atol=1e-4)
|
||||
return profiler.do_bench(backend="cupti")
|
||||
|
||||
if __name__ == "__main__":
|
||||
import csv
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--json", type=str, default=None, help="JSON file with test cases")
|
||||
parser.add_argument("--no-json", action="store_true", help="Skip JSON mode, use single-case CLI args")
|
||||
parser.add_argument("--batch", type=int, default=132, help="batch size")
|
||||
parser.add_argument("--heads", type=int, default=128, help="q heads number")
|
||||
parser.add_argument("--kv_heads", type=int, default=1, help="kv heads number")
|
||||
parser.add_argument("--kv_ctx", type=int, default=8192, help="kv context length")
|
||||
parser.add_argument("--dim", type=int, default=512, help="head dim")
|
||||
parser.add_argument("--pe_dim", type=int, default=64, help="pe head dim")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.no_json:
|
||||
main(args.batch, args.heads, args.kv_heads, args.kv_ctx, args.dim, args.pe_dim)
|
||||
else:
|
||||
json_path = args.json if args.json else "race_tests/mla/test_cases_mla_batch_ctx.json"
|
||||
with open(json_path) as f:
|
||||
cases = json.load(f)
|
||||
passed = 0
|
||||
failed = 0
|
||||
csv_path = json_path.replace(".json", "_results.csv")
|
||||
fieldnames = ["case_id", "batch", "heads", "kv_heads", "kv_ctx", "dim", "pe_dim", "latency_ms", "tflops", "status"]
|
||||
with open(csv_path, "w", newline="") as csvfile:
|
||||
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
|
||||
writer.writeheader()
|
||||
for c in cases:
|
||||
case_id = c.get("case_id", "?")
|
||||
batch = c.get("batch", 1)
|
||||
heads = c.get("heads", 16)
|
||||
kv_heads = c.get("kv_heads", 1)
|
||||
kv_ctx = c.get("kv_ctx", 8192)
|
||||
dim = c.get("dim", 512)
|
||||
pe_dim = c.get("pe_dim", 64)
|
||||
print(f"[Case {case_id}] batch={batch}, kv_ctx={kv_ctx}")
|
||||
try:
|
||||
latency, tflops = main(batch, heads, kv_heads, kv_ctx, dim, pe_dim)
|
||||
writer.writerow({
|
||||
"case_id": case_id, "batch": batch, "heads": heads, "kv_heads": kv_heads,
|
||||
"kv_ctx": kv_ctx, "dim": dim, "pe_dim": pe_dim,
|
||||
"latency_ms": round(latency, 6), "tflops": round(tflops, 4), "status": "PASS"
|
||||
})
|
||||
print(f"[Case {case_id}] PASS latency={latency:.4f}ms tflops={tflops:.4f}")
|
||||
passed += 1
|
||||
except Exception as e:
|
||||
writer.writerow({
|
||||
"case_id": case_id, "batch": batch, "heads": heads, "kv_heads": kv_heads,
|
||||
"kv_ctx": kv_ctx, "dim": dim, "pe_dim": pe_dim,
|
||||
"latency_ms": -1, "tflops": -1, "status": f"FAIL: {err_msg}"
|
||||
})
|
||||
print(f"[Case {case_id}] FAIL: {e}")
|
||||
failed += 1
|
||||
err_msg = str(e).split("\n")[0]
|
||||
print(f"\n=== Summary: {passed}/{len(cases)} passed, {failed}/{len(cases)} failed ===")
|
||||
print(f"CSV saved to: {csv_path}")
|
||||
|
||||
|
|
@ -0,0 +1,27 @@
|
|||
# MLA TileLang optimization versions
|
||||
|
||||
## Baseline
|
||||
|
||||
The tutorial implementation in `../赛题一mla教程.md` is the accepted starting
|
||||
point: `BLOCK_H=16`, `BLOCK_N=32`, `num_split=1`.
|
||||
|
||||
## Versions
|
||||
|
||||
- `opt_001_adaptive_split.py`: partitions KV across a power-of-two number of
|
||||
CTAs, targets about 128 main-kernel CTAs, keeps 16 heads together, and merges
|
||||
FP16 partial outputs using FP32 LSE and FP32 accumulation.
|
||||
- `opt_002_batch1_split_table.py`: uses an individually measured split count
|
||||
for each of the six fixed batch-one OJ shapes.
|
||||
- `run_kernel.py`: byte-identical current candidate, presently `opt_002`.
|
||||
|
||||
## Test
|
||||
|
||||
From `/data/operator_task_package/race_tests/mla`:
|
||||
|
||||
```bash
|
||||
python tilelang/test_candidate.py tilelang/opt_001_adaptive_split.py \
|
||||
--batch 1 --kv-ctx 2048
|
||||
```
|
||||
|
||||
The test utility uses PyTorch only as an out-of-kernel reference and timer. It
|
||||
is not submission code.
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,41 @@
|
|||
from test_tilelang_mla import flashattn
|
||||
|
||||
|
||||
_CACHE = {}
|
||||
|
||||
|
||||
def _choose_num_split(batch, kv_ctx):
|
||||
return 1
|
||||
|
||||
|
||||
def run_kernel(
|
||||
q,
|
||||
q_pe,
|
||||
kv,
|
||||
k_pe,
|
||||
output,
|
||||
batch,
|
||||
heads,
|
||||
kv_heads,
|
||||
kv_ctx,
|
||||
dim,
|
||||
pe_dim,
|
||||
):
|
||||
key = (batch, heads, kv_heads, kv_ctx, dim, pe_dim)
|
||||
kernel = _CACHE.get(key)
|
||||
if kernel is None:
|
||||
kernel = flashattn(
|
||||
batch,
|
||||
heads,
|
||||
kv_heads,
|
||||
kv_ctx,
|
||||
dim,
|
||||
pe_dim,
|
||||
32,
|
||||
min(16, heads // kv_heads),
|
||||
1,
|
||||
(dim + pe_dim) ** -0.5,
|
||||
)
|
||||
_CACHE[key] = kernel
|
||||
result = kernel(q, q_pe, kv, k_pe)
|
||||
output.copy_(result)
|
||||
|
|
@ -0,0 +1,169 @@
|
|||
# DeepSeek MLA Decode TileLang optimization log
|
||||
|
||||
## Baseline
|
||||
|
||||
Tutorial configuration:
|
||||
|
||||
```text
|
||||
BLOCK_H=16, BLOCK_N=32, num_split=1, threads=128
|
||||
```
|
||||
|
||||
One CTA owns all 16 heads for one batch item. The main grid therefore has only
|
||||
`batch` CTAs, while the current full C500 has 104 APs. Even `batch=32` cannot
|
||||
produce one CTA per AP.
|
||||
|
||||
Observed with the original `test_tilelang_mla.py`:
|
||||
|
||||
```text
|
||||
batch=1, kv_ctx=2048: 1.1671965 ms, 0.0610893 TFLOPS, correctness passed
|
||||
```
|
||||
|
||||
## opt_001_adaptive_split
|
||||
|
||||
Changes:
|
||||
|
||||
- Keep all 16 heads in one CTA so KV remains shared across heads.
|
||||
- Partition KV across a power-of-two `num_split`.
|
||||
- Target about 128 main-kernel CTAs without introducing KV tail tiles.
|
||||
- Merge normalized FP16 partial outputs with FP32 LSE and FP32 accumulation.
|
||||
- Main and merge kernels both use 128 threads.
|
||||
- Keep `BLOCK_N=32`, `BLOCK_H=16`, and `num_stages=0`.
|
||||
|
||||
Split selection:
|
||||
|
||||
```text
|
||||
desired_split = next_power_of_two(ceil(128 / batch))
|
||||
decrease until kv_ctx is divisible by desired_split * BLOCK_N
|
||||
```
|
||||
|
||||
### Split-path failures exposed during implementation
|
||||
|
||||
The tutorial contained a split branch, but `num_split` was fixed to one, so the
|
||||
branch had not been compiled for this shape.
|
||||
|
||||
1. `threads=256` failed MetaX MMA layout inference with integer divide by zero.
|
||||
For M=16 FullCol, the four-wave mapping produced a zero warp dimension.
|
||||
2. After changing to 128 threads, `num_stages=2` compiled but requested 92,160 B
|
||||
dynamic shared memory, exceeding the C500 65,536 B limit.
|
||||
3. The legal version uses 128 threads and `num_stages=0`.
|
||||
|
||||
### Local representative results
|
||||
|
||||
Candidate results use the same CUDA Event test utility. All use random FP16
|
||||
inputs and the OJ tolerance `rtol=2e-3, atol=2e-3`.
|
||||
|
||||
| batch | kv_ctx | split | latency ms | TFLOPS | max abs error |
|
||||
|---:|---:|---:|---:|---:|---:|
|
||||
| 1 | 2,048 | 64 | 0.025882 | 2.754975 | 0.0001221 |
|
||||
| 1 | 16,384 | 128 | 0.208020 | 2.742160 | 0.0000610 |
|
||||
| 1 | 65,536 | 128 | 0.682893 | 3.341229 | 0.0000610 |
|
||||
| 32 | 2,048 | 4 | 0.708347 | 3.221164 | 0.0002441 |
|
||||
| 32 | 65,536 | 4 | 21.359616 | 3.418341 | 0.0001221 |
|
||||
|
||||
### Same-harness comparison
|
||||
|
||||
The tutorial kernel was also measured through the same CUDA Event utility. Its
|
||||
`out_idx=[4]` profiler wrapper returns a new output tensor, so the adapter copies
|
||||
that small tensor into the supplied output. This makes the baseline comparison
|
||||
slightly conservative but does not explain the large gap.
|
||||
|
||||
| batch | kv_ctx | baseline ms | opt_001 ms | local speedup |
|
||||
|---:|---:|---:|---:|---:|
|
||||
| 1 | 2,048 | 1.178386 | 0.025882 | 45.53x |
|
||||
| 1 | 65,536 | 36.429016 | 0.682893 | 53.34x |
|
||||
| 32 | 2,048 | 1.356488 | 0.708347 | 1.92x |
|
||||
|
||||
OJ remains the final source of truth because its harness, inputs, baseline, and
|
||||
timing boundaries can differ from this local utility.
|
||||
|
||||
## Next experiments
|
||||
|
||||
1. Sweep split counts around the current policy for short, medium, and long KV.
|
||||
2. Capture mcTracer launch metadata for both main and combine kernels.
|
||||
3. Profile the long-context main kernel before changing MMA/layout.
|
||||
4. Test a single K=576 QK GEMM as an independent version.
|
||||
|
||||
## OJ result: opt_001, batch=1
|
||||
|
||||
All six batch-one cases passed on the OJ platform.
|
||||
|
||||
| case | kv_ctx | user ms | baseline ms | speedup | TFLOPS | score ratio | display |
|
||||
|---:|---:|---:|---:|---:|---:|---:|---:|
|
||||
| 1 | 2,048 | 0.027 | 0.434 | 16.074x | 2.641 | 94.40% | 94 |
|
||||
| 2 | 4,096 | 0.051 | 0.789 | 15.471x | 2.796 | 94.21% | 94 |
|
||||
| 3 | 8,192 | 0.123 | 1.541 | 12.528x | 2.319 | 92.88% | 92 |
|
||||
| 4 | 16,384 | 0.200 | 3.008 | 15.040x | 2.852 | 94.05% | 94 |
|
||||
| 5 | 32,768 | 0.351 | 5.961 | 16.983x | 3.250 | 94.74% | 94 |
|
||||
| 6 | 65,536 | 0.647 | 11.866 | 18.340x | 3.527 | 95.13% | 95 |
|
||||
|
||||
Aggregate over these six cases:
|
||||
|
||||
```text
|
||||
arithmetic-mean speedup: 15.739x
|
||||
median speedup: 15.772x
|
||||
mean score ratio: 94.235%
|
||||
mean display score: 93.833 / 100
|
||||
```
|
||||
|
||||
The OJ latency agrees with the local CUDA Event measurements. Examples:
|
||||
|
||||
```text
|
||||
b1, ctx=2,048: local 0.025882 ms, OJ 0.027 ms
|
||||
b1, ctx=16,384: local 0.208020 ms, OJ 0.200 ms
|
||||
b1, ctx=65,536: local 0.682893 ms, OJ 0.647 ms
|
||||
```
|
||||
|
||||
This agreement confirms that the local event harness measures the relevant GPU
|
||||
path and that the split/merge speedup transfers to OJ.
|
||||
|
||||
### Interpretation
|
||||
|
||||
- The tutorial baseline grows almost linearly with `kv_ctx`, as expected from a
|
||||
single CTA serially scanning every KV tile.
|
||||
- opt_001 reaches 2.3--3.5 TFLOPS and generally becomes more efficient as the
|
||||
context grows. Fixed launch, Q loading, partial-output, and combine costs are
|
||||
amortized over more KV work.
|
||||
- `kv_ctx=8,192` is the only visible dip: 2.319 TFLOPS and 12.528x speedup. Its
|
||||
neighbors use the same target of 128 CTAs, so this is not a split-count
|
||||
discontinuity. It should be repeated before changing dispatch; possible causes
|
||||
include measurement variance, clock state, or a two-tile-per-CTA scheduling
|
||||
regime.
|
||||
- The score is already saturated near 94--95. Even a large kernel speedup can
|
||||
recover at most about five display points on these cases. Optimization effort
|
||||
should first inspect the remaining batch regimes rather than overfit batch=1.
|
||||
|
||||
## opt_002_batch1_split_table
|
||||
|
||||
The OJ exposes only the six `batch=1` shapes, so opt_002 replaces the generic
|
||||
128-CTA target with a measured compile-time dispatch table:
|
||||
|
||||
| kv_ctx | opt_001 split | opt_002 split | selected local latency ms |
|
||||
|---:|---:|---:|---:|
|
||||
| 2,048 | 64 | 64 | 0.025882 |
|
||||
| 4,096 | 128 | 128 | about 0.050 |
|
||||
| 8,192 | 128 | 256 | 0.092516 |
|
||||
| 16,384 | 128 | 512 | 0.173768 |
|
||||
| 32,768 | 128 | 256 | 0.335119 |
|
||||
| 65,536 | 128 | 512 | 0.583962 |
|
||||
|
||||
Key sweep evidence:
|
||||
|
||||
```text
|
||||
ctx=8,192: split64 0.108076, split128 0.126006, split256 0.092516 ms
|
||||
ctx=16,384: split128 0.208020, split256 0.211407, split512 0.173768 ms
|
||||
ctx=65,536: split128 0.682893, split512 0.583962,
|
||||
split1024 0.729938, split2048 0.694124 ms
|
||||
```
|
||||
|
||||
The best split is not monotonic in context length. More partitions improve AP
|
||||
coverage and reduce work per CTA, but also repeat Q loads and enlarge partial
|
||||
output plus combine work. The table captures the measured balance for the fixed
|
||||
OJ shapes.
|
||||
|
||||
### Rejected companion change
|
||||
|
||||
An experimental opt_002 revision removed `acc_s_cast` and fed `S_shared`
|
||||
directly to PV GEMM. It compiled and remained correct, but the effect was not
|
||||
stable: about 1.3% faster at `ctx=8,192` and about 2% slower at `ctx=32,768` and
|
||||
`65,536`. The final opt_002 restores the opt_001 PV path and changes only split
|
||||
dispatch.
|
||||
|
|
@ -0,0 +1,272 @@
|
|||
import tilelang
|
||||
import tilelang.language as T
|
||||
|
||||
|
||||
# 固定 tile：一个 CTA 同时处理全部 16 个 query heads，并扫描一段 KV。
|
||||
BLOCK_N = 32
|
||||
BLOCK_H = 16
|
||||
# C500 有 104 个 AP，目标设置为 128 个主 kernel CTA，兼顾覆盖率和调度余量。
|
||||
TARGET_CTAS = 128
|
||||
_KERNEL_CACHE = {}
|
||||
|
||||
|
||||
def _choose_num_split(batch, kv_ctx):
|
||||
"""选择 2 的幂次 KV 分区数，并保证每个分区只包含完整的 N tile。"""
|
||||
# batch 越小，沿 KV 维切出的并行分区越多；batch 越大则减少分区合并开销。
|
||||
desired = max(1, (TARGET_CTAS + batch - 1) // batch)
|
||||
num_split = 1
|
||||
while num_split < desired:
|
||||
num_split *= 2
|
||||
|
||||
# 主循环没有 tail mask，因此 split_len 必须是 BLOCK_N 的整数倍。
|
||||
while num_split > 1 and kv_ctx % (num_split * BLOCK_N) != 0:
|
||||
num_split //= 2
|
||||
return num_split
|
||||
|
||||
|
||||
@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True})
|
||||
def flashattn_split(
|
||||
batch,
|
||||
heads,
|
||||
kv_head_num,
|
||||
seqlen_kv,
|
||||
dim,
|
||||
pe_dim,
|
||||
block_n,
|
||||
block_h,
|
||||
num_split,
|
||||
softmax_scale,
|
||||
):
|
||||
scale = float(softmax_scale * 1.44269504)
|
||||
dtype = T.float16
|
||||
accum_dtype = T.float32
|
||||
kv_group_num = heads // kv_head_num
|
||||
valid_block_h = min(block_h, kv_group_num)
|
||||
assert kv_head_num == 1, "kv_head_num must be 1"
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
Q: T.Tensor([batch, heads, dim], dtype),
|
||||
Q_pe: T.Tensor([batch, heads, pe_dim], dtype),
|
||||
KV: T.Tensor([batch, seqlen_kv, kv_head_num, dim], dtype),
|
||||
K_pe: T.Tensor([batch, seqlen_kv, kv_head_num, pe_dim], dtype),
|
||||
Output: T.Tensor([batch, heads, dim], dtype),
|
||||
):
|
||||
# 每个 split 输出一个归一化的 partial output 和对应的 log-sum-exp。
|
||||
# LSE 工作区很小，使用 FP32 可降低长上下文跨分区合并的数值误差。
|
||||
glse = T.alloc_global([batch, heads, num_split], accum_dtype)
|
||||
output_partial = T.alloc_global([batch, heads, num_split, dim], dtype)
|
||||
|
||||
# 主 kernel 网格为 batch * head_group * num_split。
|
||||
# 当前 heads=BLOCK_H=16，因此 head_group=1，不会跨 CTA 重复读取同一段 KV。
|
||||
with T.Kernel(
|
||||
batch,
|
||||
heads // min(block_h, kv_group_num),
|
||||
num_split,
|
||||
threads=128,
|
||||
) as (bid, hid, split_id):
|
||||
Q_shared = T.alloc_shared([block_h, dim], dtype)
|
||||
S_shared = T.alloc_shared([block_h, block_n], dtype)
|
||||
Q_pe_shared = T.alloc_shared([block_h, pe_dim], dtype)
|
||||
KV_shared = T.alloc_shared([block_n, dim], dtype)
|
||||
K_pe_shared = T.alloc_shared([block_n, pe_dim], dtype)
|
||||
O_shared = T.alloc_shared([block_h, dim], dtype)
|
||||
acc_s = T.alloc_fragment([block_h, block_n], accum_dtype)
|
||||
acc_s_cast = T.alloc_fragment([block_h, block_n], dtype)
|
||||
acc_o = T.alloc_fragment([block_h, dim], accum_dtype)
|
||||
scores_max = T.alloc_fragment([block_h], accum_dtype)
|
||||
scores_max_prev = T.alloc_fragment([block_h], accum_dtype)
|
||||
scores_scale = T.alloc_fragment([block_h], accum_dtype)
|
||||
scores_sum = T.alloc_fragment([block_h], accum_dtype)
|
||||
logsum = T.alloc_fragment([block_h], accum_dtype)
|
||||
cur_kv_head = hid // (kv_group_num // block_h)
|
||||
|
||||
# Q/Q_PE 在整个 KV 分区循环中保持不变，只加载一次并供所有 tile 复用。
|
||||
T.use_swizzle(10)
|
||||
T.copy(
|
||||
Q[bid, hid * valid_block_h : (hid + 1) * valid_block_h, :],
|
||||
Q_shared,
|
||||
)
|
||||
T.copy(
|
||||
Q_pe[bid, hid * valid_block_h : (hid + 1) * valid_block_h, :],
|
||||
Q_pe_shared,
|
||||
)
|
||||
T.fill(acc_o, 0)
|
||||
T.fill(logsum, 0)
|
||||
T.fill(scores_max, -T.infinity(accum_dtype))
|
||||
|
||||
split_len = seqlen_kv // num_split
|
||||
loop_range = split_len // block_n
|
||||
# num_stages=2 会把实际动态 shared 推高到 92,160 B，超过 C500 的
|
||||
# 65,536 B 上限，因此当前合法实现使用单缓冲串行流水。
|
||||
for k in T.Pipelined(loop_range, num_stages=0):
|
||||
kv_start = split_len * split_id + k * block_n
|
||||
kv_end = kv_start + block_n
|
||||
T.copy(KV[bid, kv_start:kv_end, cur_kv_head, :], KV_shared)
|
||||
T.copy(K_pe[bid, kv_start:kv_end, cur_kv_head, :], K_pe_shared)
|
||||
|
||||
# score = q @ kv^T + q_pe @ k_pe^T，两次 MMA 累加到同一 FP32 fragment。
|
||||
T.clear(acc_s)
|
||||
T.gemm(
|
||||
Q_shared,
|
||||
KV_shared,
|
||||
acc_s,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullCol,
|
||||
)
|
||||
T.gemm(
|
||||
Q_pe_shared,
|
||||
K_pe_shared,
|
||||
acc_s,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullCol,
|
||||
)
|
||||
|
||||
# 数值稳定的 online softmax：更新 running max，并把旧分母和旧输出
|
||||
# 重缩放到新的 running-max 基准下。
|
||||
T.copy(scores_max, scores_max_prev)
|
||||
T.fill(scores_max, -T.infinity(accum_dtype))
|
||||
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
|
||||
for i in T.Parallel(block_h):
|
||||
scores_max[i] = T.max(scores_max[i], scores_max_prev[i])
|
||||
scores_scale[i] = T.exp2(
|
||||
scores_max_prev[i] * scale - scores_max[i] * scale
|
||||
)
|
||||
for i, j in T.Parallel(block_h, block_n):
|
||||
acc_s[i, j] = T.exp2(
|
||||
acc_s[i, j] * scale - scores_max[i] * scale
|
||||
)
|
||||
|
||||
T.reduce_sum(acc_s, scores_sum, dim=1)
|
||||
# PV MMA 接收 FP16 概率，因此通过 shared 完成 FP32 -> FP16 转换。
|
||||
T.copy(acc_s, S_shared)
|
||||
T.copy(S_shared, acc_s_cast)
|
||||
for i in T.Parallel(block_h):
|
||||
logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
|
||||
for i, j in T.Parallel(block_h, dim):
|
||||
acc_o[i, j] *= scores_scale[i]
|
||||
T.gemm(
|
||||
acc_s_cast,
|
||||
KV_shared,
|
||||
acc_o,
|
||||
policy=T.GemmWarpPolicy.FullCol,
|
||||
)
|
||||
|
||||
# 先在每个 split 内归一化 partial output，再保存其 LSE 供第二个 kernel 合并。
|
||||
for i, j in T.Parallel(block_h, dim):
|
||||
acc_o[i, j] /= logsum[i]
|
||||
for i in T.Parallel(block_h):
|
||||
logsum[i] = T.log2(logsum[i]) + scores_max[i] * scale
|
||||
|
||||
T.copy(
|
||||
logsum,
|
||||
glse[
|
||||
bid,
|
||||
hid * valid_block_h : (hid + 1) * valid_block_h,
|
||||
split_id,
|
||||
],
|
||||
)
|
||||
T.copy(acc_o, O_shared)
|
||||
T.copy(
|
||||
O_shared,
|
||||
output_partial[
|
||||
bid,
|
||||
hid * valid_block_h : (hid + 1) * valid_block_h,
|
||||
split_id,
|
||||
:,
|
||||
],
|
||||
)
|
||||
|
||||
# 合并 kernel：用各 split 的 LSE 计算全局权重，再加权求和 partial output。
|
||||
# 每个 CTA 独占一个 (batch, head) 输出，不需要原子操作。
|
||||
with T.Kernel(heads, batch, threads=128) as (hid, bid):
|
||||
partial = T.alloc_fragment([dim], dtype)
|
||||
output_accum = T.alloc_fragment([dim], accum_dtype)
|
||||
split_lse = T.alloc_var(accum_dtype)
|
||||
merged_lse = T.alloc_var(accum_dtype)
|
||||
max_lse = T.alloc_var(accum_dtype)
|
||||
split_scale = T.alloc_var(accum_dtype)
|
||||
|
||||
T.clear(merged_lse)
|
||||
T.clear(output_accum)
|
||||
max_lse = -T.infinity(accum_dtype)
|
||||
|
||||
# 第一遍求最大 LSE，第二遍计算全局 log-sum-exp。
|
||||
for k in T.serial(num_split):
|
||||
max_lse = T.max(max_lse, glse[bid, hid, k])
|
||||
for k in T.Pipelined(num_split, num_stages=1):
|
||||
split_lse = glse[bid, hid, k]
|
||||
merged_lse += T.exp2(split_lse - max_lse)
|
||||
merged_lse = T.log2(merged_lse) + max_lse
|
||||
|
||||
# 第三遍按照 exp2(split_lse - merged_lse) 合并各分区的归一化输出。
|
||||
for k in T.serial(num_split):
|
||||
for i in T.Parallel(dim):
|
||||
partial[i] = output_partial[bid, hid, k, i]
|
||||
split_lse = glse[bid, hid, k]
|
||||
split_scale = T.exp2(split_lse - merged_lse)
|
||||
for i in T.Parallel(dim):
|
||||
output_accum[i] += partial[i] * split_scale
|
||||
|
||||
for i in T.Parallel(dim):
|
||||
Output[bid, hid, i] = output_accum[i]
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def _get_kernel(batch, heads, kv_heads, kv_ctx, dim, pe_dim):
|
||||
num_split = _choose_num_split(batch, kv_ctx)
|
||||
softmax_scale = (dim + pe_dim) ** -0.5
|
||||
# shape、tile 和 split 都属于编译期常量，必须共同进入缓存键。
|
||||
key = (
|
||||
batch,
|
||||
heads,
|
||||
kv_heads,
|
||||
kv_ctx,
|
||||
dim,
|
||||
pe_dim,
|
||||
BLOCK_N,
|
||||
BLOCK_H,
|
||||
num_split,
|
||||
)
|
||||
kernel = _KERNEL_CACHE.get(key)
|
||||
if kernel is None:
|
||||
# 每种静态规格只 JIT 一次，后续 warmup/计时直接复用已编译 kernel。
|
||||
kernel = flashattn_split(
|
||||
batch,
|
||||
heads,
|
||||
kv_heads,
|
||||
kv_ctx,
|
||||
dim,
|
||||
pe_dim,
|
||||
BLOCK_N,
|
||||
BLOCK_H,
|
||||
num_split,
|
||||
softmax_scale,
|
||||
)
|
||||
_KERNEL_CACHE[key] = kernel
|
||||
return kernel
|
||||
|
||||
|
||||
def run_kernel(
|
||||
q,
|
||||
q_pe,
|
||||
kv,
|
||||
k_pe,
|
||||
output,
|
||||
batch,
|
||||
heads,
|
||||
kv_heads,
|
||||
kv_ctx,
|
||||
dim,
|
||||
pe_dim,
|
||||
):
|
||||
kernel = _get_kernel(
|
||||
int(batch),
|
||||
int(heads),
|
||||
int(kv_heads),
|
||||
int(kv_ctx),
|
||||
int(dim),
|
||||
int(pe_dim),
|
||||
)
|
||||
kernel(q, q_pe, kv, k_pe, output)
|
||||
|
|
@ -0,0 +1,284 @@
|
|||
import tilelang
|
||||
import tilelang.language as T
|
||||
|
||||
|
||||
# 固定 tile：一个 CTA 同时处理全部 16 个 query heads，并扫描一段 KV。
|
||||
BLOCK_N = 32
|
||||
BLOCK_H = 16
|
||||
# batch=1 是当前 OJ 的完整范围；每个长度使用本地扫描得到的最佳 split。
|
||||
BATCH1_SPLITS = {
|
||||
2048: 64,
|
||||
4096: 128,
|
||||
8192: 256,
|
||||
16384: 512,
|
||||
32768: 256,
|
||||
65536: 512,
|
||||
}
|
||||
DEFAULT_TARGET_CTAS = 128
|
||||
_KERNEL_CACHE = {}
|
||||
|
||||
|
||||
def _choose_num_split(batch, kv_ctx):
|
||||
"""选择 2 的幂次 KV 分区数，并保证每个分区只包含完整的 N tile。"""
|
||||
if batch == 1 and kv_ctx in BATCH1_SPLITS:
|
||||
return BATCH1_SPLITS[kv_ctx]
|
||||
|
||||
# 保留通用回退，避免接口收到题面范围以外的规格时失效。
|
||||
# batch 越小，沿 KV 维切出的并行分区越多；batch 越大则减少分区合并开销。
|
||||
desired = max(1, (DEFAULT_TARGET_CTAS + batch - 1) // batch)
|
||||
num_split = 1
|
||||
while num_split < desired:
|
||||
num_split *= 2
|
||||
|
||||
# 主循环没有 tail mask，因此 split_len 必须是 BLOCK_N 的整数倍。
|
||||
while num_split > 1 and kv_ctx % (num_split * BLOCK_N) != 0:
|
||||
num_split //= 2
|
||||
return num_split
|
||||
|
||||
|
||||
@tilelang.jit(pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True})
|
||||
def flashattn_split(
|
||||
batch,
|
||||
heads,
|
||||
kv_head_num,
|
||||
seqlen_kv,
|
||||
dim,
|
||||
pe_dim,
|
||||
block_n,
|
||||
block_h,
|
||||
num_split,
|
||||
softmax_scale,
|
||||
):
|
||||
scale = float(softmax_scale * 1.44269504)
|
||||
dtype = T.float16
|
||||
accum_dtype = T.float32
|
||||
kv_group_num = heads // kv_head_num
|
||||
valid_block_h = min(block_h, kv_group_num)
|
||||
assert kv_head_num == 1, "kv_head_num must be 1"
|
||||
|
||||
@T.prim_func
|
||||
def main(
|
||||
Q: T.Tensor([batch, heads, dim], dtype),
|
||||
Q_pe: T.Tensor([batch, heads, pe_dim], dtype),
|
||||
KV: T.Tensor([batch, seqlen_kv, kv_head_num, dim], dtype),
|
||||
K_pe: T.Tensor([batch, seqlen_kv, kv_head_num, pe_dim], dtype),
|
||||
Output: T.Tensor([batch, heads, dim], dtype),
|
||||
):
|
||||
# 每个 split 输出一个归一化的 partial output 和对应的 log-sum-exp。
|
||||
# LSE 工作区很小，使用 FP32 可降低长上下文跨分区合并的数值误差。
|
||||
glse = T.alloc_global([batch, heads, num_split], accum_dtype)
|
||||
output_partial = T.alloc_global([batch, heads, num_split, dim], dtype)
|
||||
|
||||
# 主 kernel 网格为 batch * head_group * num_split。
|
||||
# 当前 heads=BLOCK_H=16，因此 head_group=1，不会跨 CTA 重复读取同一段 KV。
|
||||
with T.Kernel(
|
||||
batch,
|
||||
heads // min(block_h, kv_group_num),
|
||||
num_split,
|
||||
threads=128,
|
||||
) as (bid, hid, split_id):
|
||||
Q_shared = T.alloc_shared([block_h, dim], dtype)
|
||||
S_shared = T.alloc_shared([block_h, block_n], dtype)
|
||||
Q_pe_shared = T.alloc_shared([block_h, pe_dim], dtype)
|
||||
KV_shared = T.alloc_shared([block_n, dim], dtype)
|
||||
K_pe_shared = T.alloc_shared([block_n, pe_dim], dtype)
|
||||
O_shared = T.alloc_shared([block_h, dim], dtype)
|
||||
acc_s = T.alloc_fragment([block_h, block_n], accum_dtype)
|
||||
acc_s_cast = T.alloc_fragment([block_h, block_n], dtype)
|
||||
acc_o = T.alloc_fragment([block_h, dim], accum_dtype)
|
||||
scores_max = T.alloc_fragment([block_h], accum_dtype)
|
||||
scores_max_prev = T.alloc_fragment([block_h], accum_dtype)
|
||||
scores_scale = T.alloc_fragment([block_h], accum_dtype)
|
||||
scores_sum = T.alloc_fragment([block_h], accum_dtype)
|
||||
logsum = T.alloc_fragment([block_h], accum_dtype)
|
||||
cur_kv_head = hid // (kv_group_num // block_h)
|
||||
|
||||
# Q/Q_PE 在整个 KV 分区循环中保持不变，只加载一次并供所有 tile 复用。
|
||||
T.use_swizzle(10)
|
||||
T.copy(
|
||||
Q[bid, hid * valid_block_h : (hid + 1) * valid_block_h, :],
|
||||
Q_shared,
|
||||
)
|
||||
T.copy(
|
||||
Q_pe[bid, hid * valid_block_h : (hid + 1) * valid_block_h, :],
|
||||
Q_pe_shared,
|
||||
)
|
||||
T.fill(acc_o, 0)
|
||||
T.fill(logsum, 0)
|
||||
T.fill(scores_max, -T.infinity(accum_dtype))
|
||||
|
||||
split_len = seqlen_kv // num_split
|
||||
loop_range = split_len // block_n
|
||||
# num_stages=2 会把实际动态 shared 推高到 92,160 B，超过 C500 的
|
||||
# 65,536 B 上限，因此当前合法实现使用单缓冲串行流水。
|
||||
for k in T.Pipelined(loop_range, num_stages=0):
|
||||
kv_start = split_len * split_id + k * block_n
|
||||
kv_end = kv_start + block_n
|
||||
T.copy(KV[bid, kv_start:kv_end, cur_kv_head, :], KV_shared)
|
||||
T.copy(K_pe[bid, kv_start:kv_end, cur_kv_head, :], K_pe_shared)
|
||||
|
||||
# score = q @ kv^T + q_pe @ k_pe^T，两次 MMA 累加到同一 FP32 fragment。
|
||||
T.clear(acc_s)
|
||||
T.gemm(
|
||||
Q_shared,
|
||||
KV_shared,
|
||||
acc_s,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullCol,
|
||||
)
|
||||
T.gemm(
|
||||
Q_pe_shared,
|
||||
K_pe_shared,
|
||||
acc_s,
|
||||
transpose_B=True,
|
||||
policy=T.GemmWarpPolicy.FullCol,
|
||||
)
|
||||
|
||||
# 数值稳定的 online softmax：更新 running max，并把旧分母和旧输出
|
||||
# 重缩放到新的 running-max 基准下。
|
||||
T.copy(scores_max, scores_max_prev)
|
||||
T.fill(scores_max, -T.infinity(accum_dtype))
|
||||
T.reduce_max(acc_s, scores_max, dim=1, clear=False)
|
||||
for i in T.Parallel(block_h):
|
||||
scores_max[i] = T.max(scores_max[i], scores_max_prev[i])
|
||||
scores_scale[i] = T.exp2(
|
||||
scores_max_prev[i] * scale - scores_max[i] * scale
|
||||
)
|
||||
for i, j in T.Parallel(block_h, block_n):
|
||||
acc_s[i, j] = T.exp2(
|
||||
acc_s[i, j] * scale - scores_max[i] * scale
|
||||
)
|
||||
|
||||
T.reduce_sum(acc_s, scores_sum, dim=1)
|
||||
# PV MMA 接收 FP16 概率，因此通过 shared 完成 FP32 -> FP16 转换。
|
||||
T.copy(acc_s, S_shared)
|
||||
T.copy(S_shared, acc_s_cast)
|
||||
for i in T.Parallel(block_h):
|
||||
logsum[i] = logsum[i] * scores_scale[i] + scores_sum[i]
|
||||
for i, j in T.Parallel(block_h, dim):
|
||||
acc_o[i, j] *= scores_scale[i]
|
||||
T.gemm(
|
||||
acc_s_cast,
|
||||
KV_shared,
|
||||
acc_o,
|
||||
policy=T.GemmWarpPolicy.FullCol,
|
||||
)
|
||||
|
||||
# 先在每个 split 内归一化 partial output，再保存其 LSE 供第二个 kernel 合并。
|
||||
for i, j in T.Parallel(block_h, dim):
|
||||
acc_o[i, j] /= logsum[i]
|
||||
for i in T.Parallel(block_h):
|
||||
logsum[i] = T.log2(logsum[i]) + scores_max[i] * scale
|
||||
|
||||
T.copy(
|
||||
logsum,
|
||||
glse[
|
||||
bid,
|
||||
hid * valid_block_h : (hid + 1) * valid_block_h,
|
||||
split_id,
|
||||
],
|
||||
)
|
||||
T.copy(acc_o, O_shared)
|
||||
T.copy(
|
||||
O_shared,
|
||||
output_partial[
|
||||
bid,
|
||||
hid * valid_block_h : (hid + 1) * valid_block_h,
|
||||
split_id,
|
||||
:,
|
||||
],
|
||||
)
|
||||
|
||||
# 合并 kernel：用各 split 的 LSE 计算全局权重，再加权求和 partial output。
|
||||
# 每个 CTA 独占一个 (batch, head) 输出，不需要原子操作。
|
||||
with T.Kernel(heads, batch, threads=128) as (hid, bid):
|
||||
partial = T.alloc_fragment([dim], dtype)
|
||||
output_accum = T.alloc_fragment([dim], accum_dtype)
|
||||
split_lse = T.alloc_var(accum_dtype)
|
||||
merged_lse = T.alloc_var(accum_dtype)
|
||||
max_lse = T.alloc_var(accum_dtype)
|
||||
split_scale = T.alloc_var(accum_dtype)
|
||||
|
||||
T.clear(merged_lse)
|
||||
T.clear(output_accum)
|
||||
max_lse = -T.infinity(accum_dtype)
|
||||
|
||||
# 第一遍求最大 LSE，第二遍计算全局 log-sum-exp。
|
||||
for k in T.serial(num_split):
|
||||
max_lse = T.max(max_lse, glse[bid, hid, k])
|
||||
for k in T.Pipelined(num_split, num_stages=1):
|
||||
split_lse = glse[bid, hid, k]
|
||||
merged_lse += T.exp2(split_lse - max_lse)
|
||||
merged_lse = T.log2(merged_lse) + max_lse
|
||||
|
||||
# 第三遍按照 exp2(split_lse - merged_lse) 合并各分区的归一化输出。
|
||||
for k in T.serial(num_split):
|
||||
for i in T.Parallel(dim):
|
||||
partial[i] = output_partial[bid, hid, k, i]
|
||||
split_lse = glse[bid, hid, k]
|
||||
split_scale = T.exp2(split_lse - merged_lse)
|
||||
for i in T.Parallel(dim):
|
||||
output_accum[i] += partial[i] * split_scale
|
||||
|
||||
for i in T.Parallel(dim):
|
||||
Output[bid, hid, i] = output_accum[i]
|
||||
|
||||
return main
|
||||
|
||||
|
||||
def _get_kernel(batch, heads, kv_heads, kv_ctx, dim, pe_dim):
|
||||
num_split = _choose_num_split(batch, kv_ctx)
|
||||
softmax_scale = (dim + pe_dim) ** -0.5
|
||||
# shape、tile 和 split 都属于编译期常量，必须共同进入缓存键。
|
||||
key = (
|
||||
batch,
|
||||
heads,
|
||||
kv_heads,
|
||||
kv_ctx,
|
||||
dim,
|
||||
pe_dim,
|
||||
BLOCK_N,
|
||||
BLOCK_H,
|
||||
num_split,
|
||||
)
|
||||
kernel = _KERNEL_CACHE.get(key)
|
||||
if kernel is None:
|
||||
# 每种静态规格只 JIT 一次，后续 warmup/计时直接复用已编译 kernel。
|
||||
kernel = flashattn_split(
|
||||
batch,
|
||||
heads,
|
||||
kv_heads,
|
||||
kv_ctx,
|
||||
dim,
|
||||
pe_dim,
|
||||
BLOCK_N,
|
||||
BLOCK_H,
|
||||
num_split,
|
||||
softmax_scale,
|
||||
)
|
||||
_KERNEL_CACHE[key] = kernel
|
||||
return kernel
|
||||
|
||||
|
||||
def run_kernel(
|
||||
q,
|
||||
q_pe,
|
||||
kv,
|
||||
k_pe,
|
||||
output,
|
||||
batch,
|
||||
heads,
|
||||
kv_heads,
|
||||
kv_ctx,
|
||||
dim,
|
||||
pe_dim,
|
||||
):
|
||||
kernel = _get_kernel(
|
||||
int(batch),
|
||||
int(heads),
|
||||
int(kv_heads),
|
||||
int(kv_ctx),
|
||||
int(dim),
|
||||
int(pe_dim),
|
||||
)
|
||||
kernel(q, q_pe, kv, k_pe, output)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue