[Question] Performance Regression: Decode Kernel Recompiles at Every Step

[drewjin]

Required prerequisites

• I have read the documentation https://tilelang.com.
• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

Questions

Issue Description

After performing auto-tuning for the prefill kernel, I applied the selected Tile configurations (BLOCK_M, BLOCK_N, NUM_STAGES, NUM_THREADS) directly to the decode operator. However, the decode kernel triggers a re-compilation at every single step, consuming several seconds per iteration. This significantly bottlenecks the decode throughput. (Note: I haven't encountered this behavior with Triton kernels, so I am seeking guidance on the proper fix within TileLang.)

Environment

• Test Script: examples/test_sdar_diffulex_gsm8k.py
• Model: SDAR-1.7B-Chat-b32
• Decoding Strategy: block_diffusion
• Log File: log/test_sdar_diffulex_gsm8k.log

Steps to Reproduce

1. Run the test script

Setup the project Diffulex.

uv sync

Download the model JetLM/SDAR-1.7B-Chat-b32, update the model path in the script, and run:

python examples/test_sdar_diffulex_gsm8k.py 2>&1 | tee log/test_sdar_diffulex_gsm8k.log

2. Observe the logs

During the generation phase (after the first 10 outputs), observe that the decode kernel recompiles at every step.

Log Analysis

The logs in log/test_sdar_diffulex_gsm8k.log show the following:

Generating: 100%|██████████| 10/10 [02:25<00:00, 14.55s/it, Prefill=13768tok/s, Decode=42tok/s]
2025-12-25 08:28:18  [TileLang:tilelang.jit.kernel:INFO]: TileLang begins to compile kernel `kernel` with `out_idx=[10]`
2025-12-25 08:28:22  [TileLang:tilelang.jit.kernel:INFO]: TileLang completes to compile kernel `kernel`
2025-12-25 08:28:25  [TileLang:tilelang.jit.kernel:INFO]: TileLang begins to compile kernel `kernel` with `out_idx=[10]`
2025-12-25 08:28:29  [TileLang:tilelang.jit.kernel:INFO]: TileLang completes to compile kernel `kernel`
...
(Repeated recompilation, taking ~4-5 seconds per step)

**Performance Metrics:**
- **Prefill Throughput:** 13768.71 tok/s
- **Decode Throughput:** 42.53 tok/s (Severely degraded)
- **Average TPS:** 8.11 tok/s
- **Total Execution Time:** 145.53s

Code Analysis

In diffulex_kernel/python/dllm_flash_attn.py:

1. Prefill Kernel: Uses the @tilelang.autotune decorator (line 24) to perform tuning during the warmup phase. https://github.com/zhijie-group/Diffulex/blob/main/diffulex_kernel/python/dllm_flash_attn.py#L167

2. Decode Kernel: Only uses the @tilelang.jit decorator (line 164) without autotuning:
   https://github.com/zhijie-group/Diffulex/blob/main/diffulex_kernel/python/dllm_flash_attn.py#L28

   @tilelang.jit(
       out_idx=[-1], 
       pass_configs={tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,},
   )
   def dllm_flash_attn_decode_kernel(...):

3. Trigger Point: Inside the dllm_flash_attn_decode function (lines 609-632), a new kernel instance is instantiated upon every call:
   https://github.com/zhijie-group/Diffulex/blob/main/diffulex/attention/metadata.py#L20

   def dllm_flash_attn_decode(...):
       if attn_metadata.decode_mode == "static":
           # New kernel instance created per call, triggering JIT
           decode_kernel = dllm_flash_attn_decode_kernel(
               attn_metadata.num_seqs,
               ...
               **kernel_config
           )
           return decode_kernel(...)

Root Cause Analysis

1. The dllm_flash_attn_decode function instantiates a new kernel object on every invocation.
2. Even though kernel_config is passed, the out_idx=[-1] requires dynamic shape inference based on runtime tensors (logged as out_idx=[10]).
3. The TileLang JIT compiler appears to recompile for every new kernel instance and parameter combination, even if the configurations are identical.
4. Lack of a kernel instance caching mechanism leads to massive compilation overhead during the decode phase.

Expected Behavior

1. The Decode kernel should compile only once during the initial call.
2. Subsequent calls with the same parameter combinations should reuse the cached, pre-compiled kernel.
3. Eliminate the multi-second compilation latency at each decoding step.

[LeiWang1999]

@drewjin Thanks for your reprot! It's likely a bug because autotune used to have a cache.

[drewjin]

@LeiWang1999 Sorry for taking your time, there are something more I want to report under this issue.

The decoding kernel is not working on Hopper GPUs (H200).

On my Cluster_RTX_3090, the test all passed.

However, on my Cluster_H200, the process keeps hanging on this line of code, in the first test case.

(diffulex) ➜  Diffulex git:(v0.0.1) cd test/python/kernel && python test_dllm_flash_attn_decode_kernel.py
==================================================== test session starts ====================================================
platform linux -- Python 3.12.12, pytest-9.0.2, pluggy-1.6.0
rootdir: /inspire/ssd/project/advanced-machine-learning-and-deep-learning-applications/yangyi-253108120173/ssd/jyj/Diffulex
configfile: pyproject.toml
plugins: anyio-4.12.0
collected 10 items                                                                                                          

test_dllm_flash_attn_decode_kernel.py 

On NVITOP, the GPU-Util stays at MAX:

(diffulex) ➜  Diffulex git:(v0.0.1) ✗ nvitop    
Thu Dec 25 17:23:07 2025
╒═════════════════════════════════════════════════════════════════════════════╕
│ NVITOP 1.6.1      Driver Version: 570.124.06      CUDA Driver Version: 12.8 │
├───────────────────────────────┬──────────────────────┬──────────────────────┤
│ GPU  Name        Persistence-M│ Bus-Id        Disp.A │ MIG M.   Uncorr. ECC │
│ Fan  Temp  Perf  Pwr:Usage/Cap│         Memory-Usage │ GPU-Util  Compute M. │
╞═══════════════════════════════╪══════════════════════╪══════════════════════╪═════════════════════╤═══════════════════════╕
│   0  H200                  On │ 00000000:BB:00.0 Off │ Disabled           0 │ MEM: ▏ 619.1MiB     │ MBW: ▏ 0%   @ 3201MHz │
│ N/A   30C    P0   118W / 700W │  619.1MiB / 140.4GiB │    100%      Default │ UTL: ██████████ MAX │ PWR: ██ 118W    16.8% │
├───────────────────────────────┼──────────────────────┼──────────────────────┼─────────────────────┼───────────────────────┤
│   1  H200                  On │ 00000000:DB:00.0 Off │ Disabled           0 │ MEM: ▏ 3.50MiB      │ MBW: ▏ 0%   @ 3201MHz │
│ N/A   28C    P0   119W / 700W │   3.50MiB / 140.4GiB │      0%      Default │ UTL: ▏ 0%           │ PWR: ██ 119W    17.0% │
╘═══════════════════════════════╧══════════════════════╧══════════════════════╧═════════════════════╧═══════════════════════╛
[ CPU: ███▍ 4.5%                                                     UPTIME: 36.3 days ]  ( Load Average: 11.17 11.12 10.61 )
[ MEM: █████▍ 7.1%                                                      USED: 116.9GiB ]  [ SWP: ▏ 0.0%                     ]

╒═══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════╕
│ Processes:                                                                     root@diffulex-jyj--73d482b04f77-zl7upeh7ig │
│ GPU     PID      USER  GPU-MEM %SM %GMBW  %CPU  %MEM  TIME  COMMAND                                                       │
╞═══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════╡
│   0  852807 C    root 610.0MiB 100     0  99.7   0.0  1:15  python test_dllm_flash_attn_decode_kernel.py                  │
╘═══════════════════════════════════════════════════════════════════════════════════════════════════════════════════════════

[drewjin]

The generate CUDA code is as follows:

As far as I am concerned, the CUDA kernel hangs with 100% GPU utilization due to memory barrier synchronization mismatches between loading threads (threadIdx.x >= 128) and compute threads (threadIdx.x < 128). The compute threads wait for mbarrier[2] and mbarrier[3] on odd iterations (1 and 3), but the loading threads only arrive on these barriers during the final iteration (page_block_idx_local == 3), causing a deadlock in iteration 1. Additionally, when page_block_idx_global < 0, barrier arrival calls may be skipped, leading to count mismatches. The fix ensures proper barrier arrival in all iterations and page block states: mbarrier[0] and mbarrier[1] are always arrived even for invalid page blocks, and mbarrier[2]/mbarrier[3] are arrived in non-final iterations by at least one thread (via tl::tl_shuffle_elect<128>()) to satisfy the compute threads' wait conditions.

[ click to unfold / fold the code ]

#include <tl_templates/cuda/instruction/wgmma.h>
#include <math_constants.h>
#include <tl_templates/cuda/gemm.h>
#include <tl_templates/cuda/copy.h>
#include <tl_templates/cuda/reduce.h>
#include <tl_templates/cuda/ldsm.h>
#include <tl_templates/cuda/threadblock_swizzle.h>
#include <tl_templates/cuda/debug.h>
#ifdef ENABLE_BF16
#include <tl_templates/cuda/cuda_bf16_fallbacks.cuh>
#endif

extern "C" __global__ void kernel_kernel(__grid_constant__ const CUtensorMap K_Cache_desc, __grid_constant__ const CUtensorMap K_desc, bfloat16_t* __restrict__ O, __grid_constant__ const CUtensorMap Q_desc, __grid_constant__ const CUtensorMap V_Cache_desc, __grid_constant__ const CUtensorMap V_desc, const int* __restrict__ block_tables, const int* __restrict__ context_lens, const int* __restrict__ cu_seqlens_k, const int* __restrict__ cu_seqlens_q);
extern "C" __global__ void __launch_bounds__(256, 1) kernel_kernel(__grid_constant__ const CUtensorMap K_Cache_desc, __grid_constant__ const CUtensorMap K_desc, bfloat16_t* __restrict__ O, __grid_constant__ const CUtensorMap Q_desc, __grid_constant__ const CUtensorMap V_Cache_desc, __grid_constant__ const CUtensorMap V_desc, const int* __restrict__ block_tables, const int* __restrict__ context_lens, const int* __restrict__ cu_seqlens_k, const int* __restrict__ cu_seqlens_q) {
  extern __shared__ __align__(1024) uchar buf_dyn_shmem[];
  int block_table[4];
  float acc_output[64];
  float acc_score_kv[32];
  float acc_score_kvcache[16];
  float log_sum[2];
  float scores_max[2];
  float scores_max_prev[2];
  float scores_scale[2];
  float scores_sum[2];
  bfloat16_t acc_score_kvcache_cast[16];
  bfloat16_t acc_score_kv_cast[32];
  tl::GmmaDescriptor desc_a;
  tl::GmmaDescriptor desc_b;
  tl::GmmaDescriptor desc_b_1;
  tl::GmmaDescriptor desc_a_1;
  tl::GmmaDescriptor desc_b_2;
  tl::GmmaDescriptor desc_b_3;
  __shared__ uint64_t mbarrier_mem[9];
  auto mbarrier = reinterpret_cast<Barrier*>(mbarrier_mem);
  if (tl::tl_shuffle_elect<0>()) {
    tl::prefetch_tma_descriptor(Q_desc);
    tl::prefetch_tma_descriptor(K_Cache_desc);
    tl::prefetch_tma_descriptor(V_Cache_desc);
    tl::prefetch_tma_descriptor(K_desc);
    tl::prefetch_tma_descriptor(V_desc);
    mbarrier[0].init(128);
    mbarrier[1].init(128);
    mbarrier[2].init(1);
    mbarrier[3].init(1);
    mbarrier[4].init(128);
    mbarrier[5].init(128);
    mbarrier[6].init(128);
    mbarrier[7].init(128);
    mbarrier[8].init(1);
  }
  tl::fence_barrier_init();
  __syncthreads();
  if (128 <= ((int)threadIdx.x)) {
    int q_start_idx = cu_seqlens_q[0];
    int kv_start_idx = cu_seqlens_k[0];
    *(int4*)(block_table + 0) = *(int4*)(block_tables + 0);
    if (tl::tl_shuffle_elect<128>()) {
      mbarrier[8].arrive_and_expect_tx(16384);
      tl::fence_proxy_async();
      tl::tma_load(Q_desc, mbarrier[8], (&(((bfloat16_t*)buf_dyn_shmem)[0])), 0, ((int)blockIdx.y), q_start_idx);
      tl::tma_load(Q_desc, mbarrier[8], (&(((bfloat16_t*)buf_dyn_shmem)[4096])), 64, ((int)blockIdx.y), q_start_idx);
    }
    for (int page_block_idx_local = 0; page_block_idx_local < 4; ++page_block_idx_local) {
      int page_block_idx_global = block_table[page_block_idx_local];
      mbarrier[4].wait(((page_block_idx_local & 1) ^ 1));
      if (0 <= page_block_idx_global) {
        if (tl::tl_shuffle_elect<128>()) {
          mbarrier[0].expect_transaction(8192);
          tl::fence_proxy_async();
          tl::tma_load(K_Cache_desc, mbarrier[0], (&(((bfloat16_t*)buf_dyn_shmem)[28672])), 0, (((int)blockIdx.y) >> 2), 0, page_block_idx_global);
          tl::tma_load(K_Cache_desc, mbarrier[0], (&(((bfloat16_t*)buf_dyn_shmem)[30720])), 64, (((int)blockIdx.y) >> 2), 0, page_block_idx_global);
        }
      }
      mbarrier[0].arrive();
      mbarrier[5].wait(((page_block_idx_local & 1) ^ 1));
      if (0 <= page_block_idx_global) {
        if (tl::tl_shuffle_elect<128>()) {
          mbarrier[1].expect_transaction(8192);
          tl::fence_proxy_async();
          tl::tma_load(V_Cache_desc, mbarrier[1], (&(((bfloat16_t*)buf_dyn_shmem)[24576])), 0, (((int)blockIdx.y) >> 2), 0, page_block_idx_global);
          tl::tma_load(V_Cache_desc, mbarrier[1], (&(((bfloat16_t*)buf_dyn_shmem)[26624])), 64, (((int)blockIdx.y) >> 2), 0, page_block_idx_global);
        }
      }
      mbarrier[1].arrive();
      mbarrier[6].wait(((page_block_idx_local & 1) ^ 1));
      if (page_block_idx_local == 3) {
        mbarrier[2].arrive_and_expect_tx(16384);
        tl::fence_proxy_async();
        if (tl::tl_shuffle_elect<128>()) {
          tl::tma_load(K_desc, mbarrier[2], (&(((bfloat16_t*)buf_dyn_shmem)[8192])), 0, (((int)blockIdx.y) >> 2), kv_start_idx);
          tl::tma_load(K_desc, mbarrier[2], (&(((bfloat16_t*)buf_dyn_shmem)[12288])), 64, (((int)blockIdx.y) >> 2), kv_start_idx);
        }
      }
      mbarrier[7].wait(((page_block_idx_local & 1) ^ 1));
      if (page_block_idx_local == 3) {
        mbarrier[3].arrive_and_expect_tx(16384);
        tl::fence_proxy_async();
        if (tl::tl_shuffle_elect<128>()) {
          tl::tma_load(V_desc, mbarrier[3], (&(((bfloat16_t*)buf_dyn_shmem)[16384])), 0, (((int)blockIdx.y) >> 2), kv_start_idx);
          tl::tma_load(V_desc, mbarrier[3], (&(((bfloat16_t*)buf_dyn_shmem)[20480])), 64, (((int)blockIdx.y) >> 2), kv_start_idx);
        }
      }
    }
  } else {
    int q_start_idx_1 = cu_seqlens_q[0];
    int kv_start_idx_1 = cu_seqlens_k[0];
    int q_end_idx = cu_seqlens_q[1];
    int kv_end_idx = cu_seqlens_k[1];
    int cur_context_len = context_lens[0];
    *(int4*)(block_table + 0) = *(int4*)(block_tables + 0);
    #pragma unroll
    for (int i = 0; i < 32; ++i) {
      *(float2*)(acc_output + (i * 2)) = make_float2(0x0p+0f/*0.000000e+00*/, 0x0p+0f/*0.000000e+00*/);
    }
    #pragma unroll
    for (int i_1 = 0; i_1 < 16; ++i_1) {
      *(float2*)(acc_score_kv + (i_1 * 2)) = make_float2(0x0p+0f/*0.000000e+00*/, 0x0p+0f/*0.000000e+00*/);
    }
    #pragma unroll
    for (int i_2 = 0; i_2 < 8; ++i_2) {
      *(float2*)(acc_score_kvcache + (i_2 * 2)) = make_float2(0x0p+0f/*0.000000e+00*/, 0x0p+0f/*0.000000e+00*/);
    }
    #pragma unroll
    for (int i_3 = 0; i_3 < 2; ++i_3) {
      log_sum[i_3] = 0x0p+0f/*0.000000e+00*/;
    }
    #pragma unroll
    for (int i_4 = 0; i_4 < 2; ++i_4) {
      scores_max[i_4] = -CUDART_INF_F;
    }
    mbarrier[8].wait(0);
    for (int page_block_idx_local_1 = 0; page_block_idx_local_1 < 4; ++page_block_idx_local_1) {
      int page_block_idx_global_1 = block_table[page_block_idx_local_1];
      if (0 <= page_block_idx_global_1) {
        #pragma unroll
        for (int i_5 = 0; i_5 < 16; ++i_5) {
          float condval;
          if ((((q_end_idx - q_start_idx_1) <= ((((((int)threadIdx.x) >> 5) * 16) + (((i_5 & 3) >> 1) * 8)) + ((((int)threadIdx.x) & 31) >> 2))) || (cur_context_len <= ((((page_block_idx_local_1 * 32) + ((i_5 >> 2) * 8)) + ((((int)threadIdx.x) & 3) * 2)) + (i_5 & 1))))) {
            condval = -0x1.dcd65p+29f/*-1.000000e+09*/;
          } else {
            condval = 0x0p+0f/*0.000000e+00*/;
          }
          acc_score_kvcache[i_5] = condval;
        }
      }
      mbarrier[0].wait((page_block_idx_local_1 & 1));
      if (0 <= page_block_idx_global_1) {
        tl::initialize_wgmma_descriptor<1, 1, 64>(desc_a, (&(((bfloat16_t*)buf_dyn_shmem)[0])));
        tl::initialize_wgmma_descriptor<1, 1, 64>(desc_b, (&(((bfloat16_t*)buf_dyn_shmem)[28672])));
        tl::warpgroup_fence_operand(reinterpret_cast<float*>(acc_score_kvcache + 0), 16);
        tl::warpgroup_arrive();
        tl::fence_proxy_async();
        #pragma unroll
        for (int ki = 0; ki < 8; ++ki) {
          tl::wgmma_ss<tl::DataType::kBFloat16, tl::DataType::kBFloat16, tl::DataType::kFloat32, 64, 32, 16, false, false, 1, 1>(uint64_t(desc_a + ((((ki >> 2) * 8192) + ((ki & 3) * 32)) >> 4)), uint64_t(desc_b + ((((ki >> 2) * 4096) + ((ki & 3) * 32)) >> 4)), ((uint32_t*)(acc_score_kvcache + 0)), 1);
        }
        tl::warpgroup_commit_batch();
        tl::warpgroup_wait<0>();
        tl::warpgroup_fence_operand(reinterpret_cast<float*>(acc_score_kvcache + 0), 16);
      }
      mbarrier[4].arrive();
      if (0 <= page_block_idx_global_1) {
        #pragma unroll
        for (int i_6 = 0; i_6 < 2; ++i_6) {
          scores_max_prev[i_6] = scores_max[i_6];
        }
        #pragma unroll
        for (int i_7 = 0; i_7 < 2; ++i_7) {
          scores_max[i_7] = -CUDART_INF_F;
        }
        #pragma unroll
        for (int i_8 = 0; i_8 < 2; ++i_8) {
          #pragma unroll
          for (int rv = 0; rv < 8; ++rv) {
            scores_max[i_8] = max(scores_max[i_8], acc_score_kvcache[((((rv & 3) * 4) + (i_8 * 2)) + (rv >> 2))]);
          }
          scores_max[i_8] = tl::AllReduce<tl::MaxOp, 4, 1, 0, 128>::run_hopper(scores_max[i_8]);
        }
        #pragma unroll
        for (int i_9 = 0; i_9 < 2; ++i_9) {
          scores_max[i_9] = max(scores_max[i_9], scores_max_prev[i_9]);
        }
        #pragma unroll
        for (int i_10 = 0; i_10 < 2; ++i_10) {
          scores_scale[i_10] = exp2f(((scores_max_prev[i_10] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/) - (scores_max[i_10] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/)));
        }
        #pragma unroll
        for (int i_11 = 0; i_11 < 16; ++i_11) {
          acc_score_kvcache[i_11] = exp2f(((acc_score_kvcache[i_11] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/) - (scores_max[((i_11 & 3) >> 1)] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/)));
        }
        #pragma unroll
        for (int i_12 = 0; i_12 < 2; ++i_12) {
          scores_sum[i_12] = 0x0p+0f/*0.000000e+00*/;
          #pragma unroll
          for (int rv_1 = 0; rv_1 < 8; ++rv_1) {
            scores_sum[i_12] = (scores_sum[i_12] + acc_score_kvcache[((((rv_1 & 3) * 4) + (i_12 * 2)) + (rv_1 >> 2))]);
          }
          scores_sum[i_12] = tl::AllReduce<tl::SumOp, 4, 1, 0, 128>::run_hopper(scores_sum[i_12]);
        }
        #pragma unroll
        for (int i_13 = 0; i_13 < 2; ++i_13) {
          log_sum[i_13] = ((log_sum[i_13] * scores_scale[i_13]) + scores_sum[i_13]);
        }
        #pragma unroll
        for (int i_14 = 0; i_14 < 8; ++i_14) {
          uint1 __1;
          float2 v_ = *(float2*)(acc_score_kvcache + (i_14 * 2));
          (reinterpret_cast<__nv_bfloat162*>(&__1))[0] = __float22bfloat162_rn(((float2*)(&v_))[0]);
          *(uint1*)(acc_score_kvcache_cast + (i_14 * 2)) = __1;
        }
        #pragma unroll
        for (int i_15 = 0; i_15 < 64; ++i_15) {
          acc_output[i_15] = (acc_output[i_15] * scores_scale[((i_15 & 3) >> 1)]);
        }
      }
      mbarrier[1].wait((page_block_idx_local_1 & 1));
      if (0 <= page_block_idx_global_1) {
        tl::initialize_wgmma_descriptor<1, 256, 64>(desc_b_1, (&(((bfloat16_t*)buf_dyn_shmem)[24576])));
        tl::warpgroup_fence_operand(reinterpret_cast<uint32_t*>(acc_score_kvcache_cast + 0), 8);
        tl::warpgroup_fence_operand(reinterpret_cast<float*>(acc_output + 0), 64);
        tl::warpgroup_arrive();
        tl::fence_proxy_async();
        #pragma unroll
        for (int ki_1 = 0; ki_1 < 2; ++ki_1) {
          tl::wgmma_rs<tl::DataType::kBFloat16, tl::DataType::kBFloat16, tl::DataType::kFloat32, 64, 128, 16, false, true, 1, 1>(reinterpret_cast<const uint32_t*>(acc_score_kvcache_cast + (ki_1 * 8)), uint64_t(desc_b_1 + ((ki_1 * 2048) >> 4)), reinterpret_cast<uint32_t*>(acc_output + 0), 1);
        }
        tl::warpgroup_commit_batch();
        tl::warpgroup_wait<0>();
        tl::warpgroup_fence_operand(reinterpret_cast<float*>(acc_output + 0), 64);
        tl::warpgroup_fence_operand(reinterpret_cast<uint32_t*>(acc_score_kvcache_cast + 0), 8);
      }
      mbarrier[5].arrive();
      if (page_block_idx_local_1 == 3) {
        #pragma unroll
        for (int i_16 = 0; i_16 < 32; ++i_16) {
          float condval_1;
          if ((((q_end_idx - q_start_idx_1) <= ((((((int)threadIdx.x) >> 5) * 16) + (((i_16 & 3) >> 1) * 8)) + ((((int)threadIdx.x) & 31) >> 2))) || ((kv_end_idx - kv_start_idx_1) <= ((((i_16 >> 2) * 8) + ((((int)threadIdx.x) & 3) * 2)) + (i_16 & 1))))) {
            condval_1 = -0x1.dcd65p+29f/*-1.000000e+09*/;
          } else {
            condval_1 = 0x0p+0f/*0.000000e+00*/;
          }
          acc_score_kv[i_16] = condval_1;
        }
      }
      mbarrier[2].wait((page_block_idx_local_1 & 1));
      if (page_block_idx_local_1 == 3) {
        tl::initialize_wgmma_descriptor<1, 1, 64>(desc_a_1, (&(((bfloat16_t*)buf_dyn_shmem)[0])));
        tl::initialize_wgmma_descriptor<1, 1, 64>(desc_b_2, (&(((bfloat16_t*)buf_dyn_shmem)[8192])));
        tl::warpgroup_fence_operand(reinterpret_cast<float*>(acc_score_kv + 0), 32);
        tl::warpgroup_arrive();
        tl::fence_proxy_async();
        #pragma unroll
        for (int ki_2 = 0; ki_2 < 8; ++ki_2) {
          tl::wgmma_ss<tl::DataType::kBFloat16, tl::DataType::kBFloat16, tl::DataType::kFloat32, 64, 64, 16, false, false, 1, 1>(uint64_t(desc_a_1 + ((((ki_2 >> 2) * 8192) + ((ki_2 & 3) * 32)) >> 4)), uint64_t(desc_b_2 + ((((ki_2 >> 2) * 8192) + ((ki_2 & 3) * 32)) >> 4)), ((uint32_t*)(acc_score_kv + 0)), 1);
        }
        tl::warpgroup_commit_batch();
        tl::warpgroup_wait<0>();
        tl::warpgroup_fence_operand(reinterpret_cast<float*>(acc_score_kv + 0), 32);
      }
      mbarrier[6].arrive();
      if (page_block_idx_local_1 == 3) {
        #pragma unroll
        for (int i_17 = 0; i_17 < 2; ++i_17) {
          scores_max_prev[i_17] = scores_max[i_17];
        }
        #pragma unroll
        for (int i_18 = 0; i_18 < 2; ++i_18) {
          scores_max[i_18] = -CUDART_INF_F;
        }
        #pragma unroll
        for (int i_19 = 0; i_19 < 2; ++i_19) {
          #pragma unroll
          for (int rv_2 = 0; rv_2 < 16; ++rv_2) {
            scores_max[i_19] = max(scores_max[i_19], acc_score_kv[((((rv_2 & 7) * 4) + (i_19 * 2)) + (rv_2 >> 3))]);
          }
          scores_max[i_19] = tl::AllReduce<tl::MaxOp, 4, 1, 0, 128>::run_hopper(scores_max[i_19]);
        }
        #pragma unroll
        for (int i_20 = 0; i_20 < 2; ++i_20) {
          scores_max[i_20] = max(scores_max[i_20], scores_max_prev[i_20]);
        }
        #pragma unroll
        for (int i_21 = 0; i_21 < 2; ++i_21) {
          scores_scale[i_21] = exp2f(((scores_max_prev[i_21] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/) - (scores_max[i_21] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/)));
        }
        #pragma unroll
        for (int i_22 = 0; i_22 < 32; ++i_22) {
          acc_score_kv[i_22] = exp2f(((acc_score_kv[i_22] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/) - (scores_max[((i_22 & 3) >> 1)] * 0x1.0527dbd5cafffp-3f/*1.275174e-01*/)));
        }
        #pragma unroll
        for (int i_23 = 0; i_23 < 2; ++i_23) {
          scores_sum[i_23] = 0x0p+0f/*0.000000e+00*/;
          #pragma unroll
          for (int rv_3 = 0; rv_3 < 16; ++rv_3) {
            scores_sum[i_23] = (scores_sum[i_23] + acc_score_kv[((((rv_3 & 7) * 4) + (i_23 * 2)) + (rv_3 >> 3))]);
          }
          scores_sum[i_23] = tl::AllReduce<tl::SumOp, 4, 1, 0, 128>::run_hopper(scores_sum[i_23]);
        }
        #pragma unroll
        for (int i_24 = 0; i_24 < 2; ++i_24) {
          log_sum[i_24] = ((log_sum[i_24] * scores_scale[i_24]) + scores_sum[i_24]);
        }
        #pragma unroll
        for (int i_25 = 0; i_25 < 16; ++i_25) {
          uint1 __2;
          float2 v__1 = *(float2*)(acc_score_kv + (i_25 * 2));
          (reinterpret_cast<__nv_bfloat162*>(&__2))[0] = __float22bfloat162_rn(((float2*)(&v__1))[0]);
          *(uint1*)(acc_score_kv_cast + (i_25 * 2)) = __2;
        }
        #pragma unroll
        for (int i_26 = 0; i_26 < 64; ++i_26) {
          acc_output[i_26] = (acc_output[i_26] * scores_scale[((i_26 & 3) >> 1)]);
        }
      }
      mbarrier[3].wait((page_block_idx_local_1 & 1));
      if (page_block_idx_local_1 == 3) {
        tl::initialize_wgmma_descriptor<1, 512, 64>(desc_b_3, (&(((bfloat16_t*)buf_dyn_shmem)[16384])));
        tl::warpgroup_fence_operand(reinterpret_cast<uint32_t*>(acc_score_kv_cast + 0), 16);
        tl::warpgroup_fence_operand(reinterpret_cast<float*>(acc_output + 0), 64);
        tl::warpgroup_arrive();
        tl::fence_proxy_async();
        #pragma unroll
        for (int ki_3 = 0; ki_3 < 4; ++ki_3) {
          tl::wgmma_rs<tl::DataType::kBFloat16, tl::DataType::kBFloat16, tl::DataType::kFloat32, 64, 128, 16, false, true, 1, 1>(reinterpret_cast<const uint32_t*>(acc_score_kv_cast + (ki_3 * 8)), uint64_t(desc_b_3 + ((ki_3 * 2048) >> 4)), reinterpret_cast<uint32_t*>(acc_output + 0), 1);
        }
        tl::warpgroup_commit_batch();
        tl::warpgroup_wait<0>();
        tl::warpgroup_fence_operand(reinterpret_cast<float*>(acc_output + 0), 64);
        tl::warpgroup_fence_operand(reinterpret_cast<uint32_t*>(acc_score_kv_cast + 0), 16);
      }
      mbarrier[7].arrive();
    }
    #pragma unroll
    for (int i_27 = 0; i_27 < 64; ++i_27) {
      acc_output[i_27] = (acc_output[i_27] / log_sum[((i_27 & 3) >> 1)]);
    }
    tl::__sync_thread_partial<3, 128>();
    #pragma unroll
    for (int i_28 = 0; i_28 < 8; ++i_28) {
      tl::ptx_stmatrix_x4((&(((bfloat16_t*)buf_dyn_shmem)[(((((((i_28 >> 2) * 4096) + ((((int)threadIdx.x) >> 5) * 1024)) + ((((int)threadIdx.x) & 15) * 64)) + (((((((int)threadIdx.x) & 7) >> 2) + ((i_28 & 3) >> 1)) & 1) * 32)) + (((((((int)threadIdx.x) & 3) >> 1) + (i_28 & 1)) & 1) * 16)) + (((((((int)threadIdx.x) & 31) >> 4) + (((int)threadIdx.x) & 1)) & 1) * 8))])), __pack_half2(((bfloat16_t)acc_output[(i_28 * 8)]), ((bfloat16_t)acc_output[((i_28 * 8) + 1)])), __pack_half2(((bfloat16_t)acc_output[((i_28 * 8) + 2)]), ((bfloat16_t)acc_output[((i_28 * 8) + 3)])), __pack_half2(((bfloat16_t)acc_output[((i_28 * 8) + 4)]), ((bfloat16_t)acc_output[((i_28 * 8) + 5)])), __pack_half2(((bfloat16_t)acc_output[((i_28 * 8) + 6)]), ((bfloat16_t)acc_output[((i_28 * 8) + 7)])));
    }
    tl::__sync_thread_partial<3, 128>();
    #pragma unroll
    for (int i_29 = 0; i_29 < 8; ++i_29) {
      if (((i_29 * 8) + (((int)threadIdx.x) >> 4)) < (q_end_idx - q_start_idx_1)) {
        if (0 <= (((i_29 * 8) + (((int)threadIdx.x) >> 4)) + q_start_idx_1)) {
          if (((((((int)threadIdx.x) >> 4) + q_start_idx_1) >> 3) + i_29) < 4) {
            *(uint4*)(O + (((((((int64_t)i_29) * (int64_t)32768) + ((((int64_t)((int)threadIdx.x)) >> (int64_t)4) * (int64_t)4096)) + (((int64_t)q_start_idx_1) * (int64_t)4096)) + (((int64_t)((int)blockIdx.y)) * (int64_t)128)) + ((((int64_t)((int)threadIdx.x)) & (int64_t)15) * (int64_t)8))) = *(uint4*)(((bfloat16_t*)buf_dyn_shmem) + ((((((((((int)threadIdx.x) & 15) >> 3) * 4096) + (i_29 * 512)) + ((((int)threadIdx.x) >> 4) * 64)) + ((((((int)threadIdx.x) >> 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)));
          }
        }
      }
    }
  }
}

[LeiWang1999]

@drewjin it's something related to tma lowering and auto warp specialization, we are now considering refactoring this part, for now we can use

@tilelang.jit(
    pass_configs={
        tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
        tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
        tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
    }
)

as a workaround.

[LeiWang1999]

For the duplicated autotuning issue, I suspect it's caused by treating all shapes as static. When the compiler encounters a specific shape, it triggers autotuning and compilation processes. In such cases, we can use T.dynamic to make some shapes dynamic, which should help avoid redundant autotuning.

[drewjin]

@drewjin it's something related to tma lowering and auto warp specialization, we are now considering refactoring this part, for now we can use

@tilelang.jit(
pass_configs={
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: True,
tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
}
)
as a workaround.

@LeiWang1999 Thanks for the reply, I will try on my H200 cluster later.

[drewjin]

For the duplicated autotuning issue, I suspect it's caused by treating all shapes as static. When the compiler encounters a specific shape, it triggers autotuning and compilation processes. In such cases, we can use T.dynamic to make some shapes dynamic, which should help avoid redundant autotuning.

@LeiWang1999 , thank you for your reply. I have resolved this issue by using T.dynamic. I would also like to thank @Rachmanino for his guidance on this matter yesterday on wechat.

I initially misunderstood this feature, assuming it worked similarly to constexpr in Triton, but I failed to notice the specific syntax for defining dynamic parameters in TileLang. I've since learned that in TileLang, parameters are const by default, and you must explicitly use T.dynamic when a value needs to be dynamic.