[Perf] Skip blocks past per-request live length in full-width Triton kernels (#32109)

2026-07-22 22:00:25 -07:00
parent eb242b6c03
commit 1b63155efe
9 changed files with 187 additions and 28 deletions
@@ -639,6 +639,9 @@ def _get_k_and_s_triton_kernel(
    k_offsets = thread_idx * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)

    seq_len = tl.load(seq_len_ptr + batch_id)
+    # Grid axis 1 spans the batch-max seq len; fully-masked blocks store nothing.
+    if block_token_start >= seq_len:
+        return
    token_valid_mask = token_ids < seq_len

    pre_batch_idx = tl.arange(0, seq_len_num_pow)
@@ -87,6 +87,9 @@ def transform_index_page_table_prefill_kernel(

    query_start = tl.load(cu_seqlens_q_ptr + request_id)
    query_end = tl.load(cu_seqlens_q_ptr + request_id + 1)
+    # Grid axis 1 spans the batch-max extend len; fully-masked blocks store nothing.
+    if query_start + tl.program_id(1) * BLOCK_Q >= query_end:
+        return
    token_indices = query_start + query_offsets
    mask = (token_indices[:, None] < query_end) & (topk_offsets[None, :] < TOPK)

@@ -71,6 +71,16 @@ def _fused_dsa_decode_metadata_kernel(
        mask=row < bs,
        other=0,
    )
+    # Skip column blocks past the request's kv length: no consumer reads there
+    # (attention and the indexer both stay within cache_seqlens). Loaded after
+    # req_idx so the two scalar loads pipeline (no added latency when live).
+    kv_len = tl.load(
+        seq_lens + row * seq_lens_stride,
+        mask=row < bs,
+        other=0,
+    ).to(tl.int32)
+    if col_block * BLOCK_N >= kv_len:
+        return
    vals = tl.load(
        req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1,
        mask=mask,
@@ -120,6 +130,10 @@ def fused_dsa_decode_metadata(
    where the wide table is never read (attention uses topk_indices, the indexer
    uses real_page_table); ``real_page_size`` must be >1 in that case. When a
    tensor is passed, behavior is unchanged (both tables are written).
+
+    Contract: each page-table row is written only over its live prefix
+    ([:cache_seqlens]); the tail keeps stale values across CUDA-graph replays, so
+    consumers must bound reads by cache_seqlens.
    """
    assert seq_lens.is_cuda
    assert req_pool_indices.is_cuda
@@ -283,6 +297,20 @@ def _fused_dsa_target_verify_metadata_kernel(
        mask=out_row < expanded_size,
        other=0,
    )
+    # Skip column blocks past the request's kv length (seq_len + next_n): no
+    # consumer reads there (attention and the indexer stay within cache_seqlens).
+    # Loaded after req_idx so the two scalar loads pipeline (no added latency
+    # when live).
+    kv_len = (
+        tl.load(
+            seq_lens + req_row * seq_lens_stride,
+            mask=out_row < expanded_size,
+            other=0,
+        ).to(tl.int32)
+        + next_n
+    )
+    if col_block * BLOCK_N >= kv_len:
+        return
    vals = tl.load(
        req_to_token + req_idx * req_to_token_stride_0 + offs_n * req_to_token_stride_1,
        mask=mask,
@@ -313,6 +313,10 @@ def _fwd_kernel(
    cur_seq_len_prefix = tl.load(kv_indptr + cur_seq + 1) - cur_seq_kv_start_idx
    cur_seq_len = cur_seq_len_prefix + cur_seq_len_extend

+    # Grid axis 2 spans the batch-max extend length; all stores are masked by mask_m.
+    if cur_block_m * BLOCK_M >= cur_seq_len_extend:
+        return
+
    if USE_CUSTOM_MASK:
        cur_seq_mask_start_idx = tl.load(mask_indptr + cur_seq)

@@ -911,6 +915,10 @@ def _fwd_kernel_unified(
    cur_seq_kv_len = tl.load(kv_indptr + cur_seq + 1) - cur_seq_kv_start_idx
    cur_seq_prefix_len = tl.load(prefix_lens + cur_seq)

+    # Grid axis 2 spans the batch-max extend length; the store is masked by mask_m.
+    if cur_block_m * BLOCK_M >= cur_seq_q_len:
+        return
+
    # Load window start position for sliding window attention
    # This is the absolute position of the first key in the window (0 if no sliding window)
    cur_window_start = 0
@@ -213,13 +213,25 @@ def _fused_metadata_kernel_general(
        return

    i = pid_b
+    # Self-guard on the device-side seq_len: skip column chunks past the
+    # request's live pages (tails keep stale values the attention kernels
+    # never read past cache_seqlens).
+    seq_len = tl.load(seq_lens + i * seq_lens_stride_0).to(tl.int32)
+    if page_size == 1:
+        num_live_pages = seq_len + seq_len_delta
+    else:
+        num_live_pages = (seq_len + seq_len_delta + (1 << SHIFT) - 1) >> SHIFT
+    num_live_pages = tl.minimum(num_live_pages, max_seq_pages)
+    col_start = pid_c * BLOCK_COLS
+    if col_start >= num_live_pages:
+        return
+
    # Load row index for this batch (all threads in block have same i)
    row_idx = tl.load(req_pool_indices + i * req_pool_indices_stride_0)
    row_offset = row_idx * req_to_token_stride_0

-    col_start = pid_c * BLOCK_COLS
    col_offsets = col_start + tl.arange(0, BLOCK_COLS)
-    mask = col_offsets < max_seq_pages
+    mask = col_offsets < num_live_pages

    # Compute column indices in the source tensor (token offset)
    if page_size == 1:
@@ -303,13 +315,21 @@ def _fused_metadata_kernel_ps1_no_swa(
        return

    i = pid_b
+    # Self-guard on the device-side seq_len: skip column chunks past the
+    # request's live pages (tails keep stale values the attention kernels
+    # never read past cache_seqlens).
+    seq_len = tl.load(seq_lens + i * seq_lens_stride_0).to(tl.int32)
+    num_live_pages = tl.minimum(seq_len + seq_len_delta, max_seq_pages)
+    col_start = pid_c * BLOCK_COLS
+    if col_start >= num_live_pages:
+        return
+
    # Load row index for this batch (all threads in block have same i)
    row_idx = tl.load(req_pool_indices + i * req_pool_indices_stride_0)
    row_offset = row_idx * req_to_token_stride_0

-    col_start = pid_c * BLOCK_COLS
    col_offsets = col_start + tl.arange(0, BLOCK_COLS)
-    mask = col_offsets < max_seq_pages
+    mask = col_offsets < num_live_pages

    # page_size = 1: col_idx = col_offsets
    rt_offsets = row_offset + col_offsets * req_to_token_stride_1
@@ -555,6 +575,10 @@ def normal_decode_set_metadata(
      5. (optional) swa_page_table for sliding window attention

    Achieves ~5.2x speedup on H200 hardware for typical decode workloads.
+
+    Contract: only the live prefix (cdiv(cache_seqlens, page_size) pages) of each
+    page_table / swa_page_table row is (re)written; the tail keeps stale values
+    across CUDA-graph replays, so consumers must bound reads by cache_seqlens.
    """
    assert (
        page_size > 0 and (page_size & (page_size - 1)) == 0
@@ -75,9 +75,12 @@ def update_trtllm_mha_graph_metadata_kernel(
        row_out = page_table_ptr + pid.to(tl.int64) * page_table_stride
        if HAS_SWA:
            swa_row_out = swa_page_table_ptr + pid.to(tl.int64) * swa_page_table_stride
-        for i in range(tl.cdiv(max_seq_pages, PAGE_BLOCK)):
+        # Self-guard on the device-side seqlen: pages past cdiv(cache_seqlen,
+        # PAGE_SIZE) keep stale values the attention kernels never read.
+        num_live_pages = tl.minimum(tl.cdiv(seqlen, PAGE_SIZE), max_seq_pages)
+        for i in range(tl.cdiv(num_live_pages, PAGE_BLOCK)):
            page_idx = i * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK)
-            mask = page_idx < max_seq_pages
+            mask = page_idx < num_live_pages
            token = tl.load(
                row_in + page_idx.to(tl.int64) * PAGE_SIZE, mask=mask, other=0
            )
@@ -143,7 +146,12 @@ def update_trtllm_mha_graph_metadata(
    q_stride: int = 0,
    q_mode: int = Q_MODE_NONE,
 ):
-    """Launch the fused metadata update (one kernel for the whole replay init)."""
+    """Launch the fused metadata update (one kernel for the whole replay init).
+
+    Contract: only the live prefix (cdiv(cache_seqlens, page_size) pages) of each
+    page_table / swa_page_table row is (re)written; the tail keeps stale values
+    across replays, so consumers must bound reads by cache_seqlens.
+    """
    if bs == 0:
        return