From c4ec39a785c5ef543bd61a32a5bc284a2d1abc75 Mon Sep 17 00:00:00 2001 From: amd-danli103 Date: Mon, 15 Jun 2026 18:26:59 +0800 Subject: [PATCH] [AMD] refactor sparse MLA decode kernel for Deepseek V4 triton backend (#28265) Co-authored-by: Raiden-Makoto Co-authored-by: yichiche@amd.com --- .../triton_mla_kernels_decode_common.py | 585 ------- .../triton_mla_kernels_decode_dsv4.py | 1355 ----------------- .../triton_mla_kernels_decode_fused.py | 313 ++-- .../triton_mla_kernels_decode_optimized.py | 343 +---- .../triton_mla_kernels_decode_splitk.py | 534 ------- 5 files changed, 212 insertions(+), 2918 deletions(-) delete mode 100644 python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_common.py delete mode 100644 python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_dsv4.py delete mode 100644 python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_splitk.py diff --git a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_common.py b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_common.py deleted file mode 100644 index ac9c1cced..000000000 --- a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_common.py +++ /dev/null @@ -1,585 +0,0 @@ -""" -Common utilities and attention kernels for Triton MLA Decode. - -This module contains shared code for the DeepSeek V4 Triton decode implementation: -- Attention kernels (unified sparse decode) -- Helper functions for chunked attention -- Token range computation for memory-based chunking -""" - -from typing import List, Tuple - -import torch -import triton -import triton.language as tl - -LOG2E = tl.constexpr(1.4426950408889634) - - -# ============================================================================ -# Bucketing for autotune keys to avoid recompilation per unique batch size -# ============================================================================ -def _bucket_total_tokens(total_tokens: int) -> int: - """Round total_tokens up to the nearest power of 2 for autotune key stability. - - In serving, total_tokens (= batch_size * seq_len) varies with every batch. - Using the exact value as an autotune key causes recompilation for each unique - value. Bucketing to powers of 2 limits the number of unique keys to ~15, - dramatically reducing autotuning overhead. - - Returns: - Power-of-2 bucket: 1, 2, 4, 8, ..., up to the next power of 2. - """ - if total_tokens <= 0: - return 1 - # Round up to next power of 2 - n = 1 - while n < total_tokens: - n <<= 1 - return n - - -# ============================================================================ -# Helper function to compute workload size category for autotune -# ============================================================================ -def _get_workload_size_category(total_tokens: int, topk: int) -> int: - """ - Compute workload size category for autotune key. - Returns: - 0: small (< 10K elements) - 1: medium (10K - 100K elements) - 2: large (100K - 1M elements) - 3: very large (> 1M elements) - """ - total_elements = total_tokens * topk - if total_elements < 10000: - return 0 - elif total_elements < 100000: - return 1 - elif total_elements < 1000000: - return 2 - else: - return 3 - - -# ============================================================================ -# Unified Attention Kernels -# ============================================================================ - - -# ============================================================================ -# CDNA4 (gfx950) Optimized: Added high-performance configs for MI355X -# Best config for h_q=128, large topk: BLOCK_H=64, BLOCK_N=256, num_warps=8 -# ============================================================================ -@triton.autotune( - configs=[ - # Selected based on CDNA4 architecture analysis: - # - BLOCK_D=128 is fixed (matches KV tile structure for d_qk=512). - # - BLOCK_N=256: best for amortizing memory access over topk dimension. - # (decode attention is memory-bound; larger BLOCK_N = fewer iterations) - # - num_warps=8: memory-bound decode benefits from more warps for latency hiding. - # - BLOCK_H varies to cover different batch sizes: - # * BLOCK_H=16: cdiv(128,16)=8 H-blocks, best for small batches (bs=1-8) - # * BLOCK_H=32: cdiv(128,32)=4 H-blocks, good for medium batches (bs=8-32) - # * BLOCK_H=64: cdiv(128,64)=2 H-blocks, best for large batches (bs=32+) - # (original comment: "Best for h_q=128, large topk") - # * BLOCK_H=128: cdiv(128,128)=1 H-block, for very large batches (bs=128+) - triton.Config( - {"BLOCK_H": 16, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=1 - ), - triton.Config( - {"BLOCK_H": 32, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=1 - ), - triton.Config( - {"BLOCK_H": 64, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=1 - ), - triton.Config( - {"BLOCK_H": 128, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=1 - ), - ], - key=["total_tokens_bucket", "h_q", "total_topk", "d_qk"], -) -@triton.jit -def _unified_sparse_decode_kernel( - Q, - KV, - Mask, - AttnSink, - Output, - LSE, - sm_scale, - total_tokens, - total_tokens_bucket, - h_q, - total_topk, - d_qk, - d_v, - stride_q_t, - stride_q_h, - stride_q_d, - stride_kv_t, - stride_kv_k, - stride_kv_d, - stride_mask_t, - stride_mask_k, - stride_o_t, - stride_o_h, - stride_o_d, - stride_lse_t, - stride_lse_h, - HAS_ATTN_SINK: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_D: tl.constexpr, -): - """Unified attention kernel with single KV buffer (int64 safe).""" - pid_t = tl.program_id(0) - pid_h = tl.program_id(1) - pid_t_64 = pid_t.to(tl.int64) - - NEG_INF = float("-inf") - POS_INF = float("+inf") - - offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) - mask_h = offs_h < h_q - - m_i = tl.full([BLOCK_H], NEG_INF, dtype=tl.float32) - l_i = tl.zeros([BLOCK_H], dtype=tl.float32) - - acc_0 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - acc_1 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - acc_2 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - acc_3 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - - stride_q_t_64 = tl.cast(stride_q_t, tl.int64) - stride_kv_t_64 = tl.cast(stride_kv_t, tl.int64) - stride_mask_t_64 = tl.cast(stride_mask_t, tl.int64) - q_base = Q + pid_t_64 * stride_q_t_64 - kv_base = KV + pid_t_64 * stride_kv_t_64 - mask_base = Mask + pid_t_64 * stride_mask_t_64 - - for n_start in range(0, total_topk, BLOCK_N): - offs_n = n_start + tl.arange(0, BLOCK_N) - mask_n = offs_n < total_topk - - mask_ptrs = mask_base + offs_n * stride_mask_k - invalid = tl.load(mask_ptrs, mask=mask_n, other=True) - valid = mask_n & ~invalid - - qk = tl.zeros([BLOCK_H, BLOCK_N], dtype=tl.float32) - - for d_start in range(0, d_qk, BLOCK_D): - offs_d = d_start + tl.arange(0, BLOCK_D) - mask_d = offs_d < d_qk - - q_ptrs = ( - q_base + offs_h[:, None] * stride_q_h + offs_d[None, :] * stride_q_d - ) - q_chunk = tl.load( - q_ptrs, mask=mask_h[:, None] & mask_d[None, :], other=0.0 - ).to(tl.bfloat16) - - k_ptrs = ( - kv_base + offs_n[:, None] * stride_kv_k + offs_d[None, :] * stride_kv_d - ) - k_chunk = tl.load( - k_ptrs, mask=valid[:, None] & mask_d[None, :], other=0.0 - ).to(tl.bfloat16) - - qk += tl.dot(q_chunk, tl.trans(k_chunk)) - - qk = qk * sm_scale - qk = tl.where(valid[None, :], qk, NEG_INF) - - m_ij = tl.max(qk, axis=1) - m_new = tl.maximum(m_i, m_ij) - alpha = tl.where(m_i == NEG_INF, 0.0, tl.math.exp2((m_i - m_new) * LOG2E)) - p = tl.where(qk == NEG_INF, 0.0, tl.math.exp2((qk - m_new[:, None]) * LOG2E)) - l_new = alpha * l_i + tl.sum(p, axis=1) - p_bf16 = p.to(tl.bfloat16) - - offs_v = tl.arange(0, BLOCK_D) - v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d - v = tl.load(v_ptrs, mask=valid[:, None], other=0.0).to(tl.bfloat16) - acc_0 = acc_0 * alpha[:, None] + tl.dot(p_bf16, v) - - offs_v = BLOCK_D + tl.arange(0, BLOCK_D) - v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d - v = tl.load( - v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0 - ).to(tl.bfloat16) - acc_1 = acc_1 * alpha[:, None] + tl.dot(p_bf16, v) - - offs_v = 2 * BLOCK_D + tl.arange(0, BLOCK_D) - v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d - v = tl.load( - v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0 - ).to(tl.bfloat16) - acc_2 = acc_2 * alpha[:, None] + tl.dot(p_bf16, v) - - offs_v = 3 * BLOCK_D + tl.arange(0, BLOCK_D) - v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d - v = tl.load( - v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0 - ).to(tl.bfloat16) - acc_3 = acc_3 * alpha[:, None] + tl.dot(p_bf16, v) - - m_i = m_new - l_i = l_new - - lse = m_i + tl.math.log2(tl.where(l_i == 0.0, 1.0, l_i)) / LOG2E - is_lonely_q = l_i == 0.0 - - if HAS_ATTN_SINK: - attn_sink_vals = tl.load(AttnSink + offs_h, mask=mask_h, other=0.0) - exp_attn_sink_minus_m = tl.math.exp2((attn_sink_vals - m_i) * LOG2E) - denominator = l_i + exp_attn_sink_minus_m - denominator = tl.where(denominator == 0.0, 1.0, denominator) - output_scale = 1.0 / denominator - else: - output_scale = tl.where(l_i == 0.0, 0.0, 1.0 / l_i) - - # Pre-compute 2D versions for efficiency - is_lonely_q_2d = is_lonely_q[:, None] - output_scale_2d = output_scale[:, None] - acc_0 = tl.where(is_lonely_q_2d, 0.0, acc_0 * output_scale_2d) - acc_1 = tl.where(is_lonely_q_2d, 0.0, acc_1 * output_scale_2d) - acc_2 = tl.where(is_lonely_q_2d, 0.0, acc_2 * output_scale_2d) - acc_3 = tl.where(is_lonely_q_2d, 0.0, acc_3 * output_scale_2d) - lse = tl.where(is_lonely_q, POS_INF, lse) - - stride_lse_t_64 = tl.cast(stride_lse_t, tl.int64) - tl.store(LSE + pid_t_64 * stride_lse_t_64 + offs_h * stride_lse_h, lse, mask=mask_h) - - stride_o_t_64 = tl.cast(stride_o_t, tl.int64) - o_base = Output + pid_t_64 * stride_o_t_64 - # Pre-compute 2D versions - offs_h_2d = offs_h[:, None] - mask_h_2d = mask_h[:, None] - offs_v_0 = tl.arange(0, BLOCK_D) - offs_v_1 = BLOCK_D + tl.arange(0, BLOCK_D) - offs_v_2 = 2 * BLOCK_D + tl.arange(0, BLOCK_D) - offs_v_3 = 3 * BLOCK_D + tl.arange(0, BLOCK_D) - tl.store( - o_base + offs_h_2d * stride_o_h + offs_v_0[None, :] * stride_o_d, - acc_0.to(tl.bfloat16), - mask=mask_h_2d, - ) - tl.store( - o_base + offs_h_2d * stride_o_h + offs_v_1[None, :] * stride_o_d, - acc_1.to(tl.bfloat16), - mask=mask_h_2d & (offs_v_1[None, :] < d_v), - ) - tl.store( - o_base + offs_h_2d * stride_o_h + offs_v_2[None, :] * stride_o_d, - acc_2.to(tl.bfloat16), - mask=mask_h_2d & (offs_v_2[None, :] < d_v), - ) - tl.store( - o_base + offs_h_2d * stride_o_h + offs_v_3[None, :] * stride_o_d, - acc_3.to(tl.bfloat16), - mask=mask_h_2d & (offs_v_3[None, :] < d_v), - ) - - -# ============================================================================ -# Attention Runner Functions -# ============================================================================ - - -def run_unified_attention( - q_reshaped, - gathered_kv, - invalid_mask, - d_v, - sm_scale, - total_tokens, - h_q, - total_topk, - d_qk, - attn_sink=None, -): - """Run unified attention with single KV buffer. - - Run unified sparse decode attention kernel. - """ - output = torch.empty( - (total_tokens, h_q, d_v), dtype=torch.bfloat16, device=q_reshaped.device - ) - lse = torch.empty( - (total_tokens, h_q), dtype=torch.float32, device=q_reshaped.device - ) - - HAS_ATTN_SINK = attn_sink is not None - attn_sink_tensor = attn_sink if HAS_ATTN_SINK else lse[:1] - - grid = lambda meta: (total_tokens, triton.cdiv(h_q, meta["BLOCK_H"])) - _unified_sparse_decode_kernel[grid]( - q_reshaped, - gathered_kv, - invalid_mask, - attn_sink_tensor, - output, - lse, - sm_scale, - total_tokens, - _bucket_total_tokens(total_tokens), - h_q, - total_topk, - d_qk, - d_v, - q_reshaped.stride(0), - q_reshaped.stride(1), - q_reshaped.stride(2), - gathered_kv.stride(0), - gathered_kv.stride(1), - gathered_kv.stride(2), - invalid_mask.stride(0), - invalid_mask.stride(1), - output.stride(0), - output.stride(1), - output.stride(2), - lse.stride(0), - lse.stride(1), - HAS_ATTN_SINK=HAS_ATTN_SINK, - ) - return output, lse - - -def run_chunked_attention_triton( - q_reshaped, - gathered_kv, - invalid_mask, - d_v, - sm_scale, - total_tokens, - h_q, - total_topk, - d_qk, - attn_sink=None, - chunk_size=8192, -): - """Chunked attention using Triton kernels with cross-chunk softmax merging.""" - device = q_reshaped.device - - num_chunks = (total_topk + chunk_size - 1) // chunk_size - - kv_chunks = [] - mask_chunks = [] - chunk_sizes = [] - - for chunk_idx in range(num_chunks): - start_k = chunk_idx * chunk_size - end_k = min(start_k + chunk_size, total_topk) - chunk_topk = end_k - start_k - chunk_sizes.append(chunk_topk) - kv_chunks.append(gathered_kv[:, start_k:end_k, :].contiguous()) - mask_chunks.append(invalid_mask[:, start_k:end_k].contiguous()) - - lse_acc = torch.full( - (total_tokens, h_q), float("-inf"), dtype=torch.float32, device=device - ) - acc = torch.zeros((total_tokens, h_q, d_v), dtype=torch.float32, device=device) - - for chunk_idx in range(num_chunks): - kv_chunk = kv_chunks[chunk_idx] - mask_chunk = mask_chunks[chunk_idx] - chunk_topk = chunk_sizes[chunk_idx] - - chunk_output, chunk_lse = run_unified_attention( - q_reshaped, - kv_chunk, - mask_chunk, - d_v, - sm_scale, - total_tokens, - h_q, - chunk_topk, - d_qk, - attn_sink=None, - ) - - is_chunk_lonely = torch.isinf(chunk_lse) & (chunk_lse > 0) - - chunk_lse_for_merge = torch.where( - is_chunk_lonely, torch.full_like(chunk_lse, float("-inf")), chunk_lse - ) - - lse_max = torch.maximum(lse_acc, chunk_lse_for_merge) - - exp_acc = torch.exp(lse_acc - lse_max) - exp_acc = torch.where(torch.isnan(exp_acc), torch.zeros_like(exp_acc), exp_acc) - - exp_chunk = torch.exp(chunk_lse_for_merge - lse_max) - exp_chunk = torch.where( - torch.isnan(exp_chunk) | is_chunk_lonely, - torch.zeros_like(exp_chunk), - exp_chunk, - ) - - sum_exp = exp_acc + exp_chunk - lse_new = lse_max + torch.log( - torch.where(sum_exp == 0, torch.ones_like(sum_exp), sum_exp) - ) - - both_empty = (lse_acc == float("-inf")) & (chunk_lse_for_merge == float("-inf")) - lse_new = torch.where( - both_empty, torch.full_like(lse_new, float("-inf")), lse_new - ) - - weight_acc = torch.exp(lse_acc - lse_new) - weight_acc = torch.where( - torch.isnan(weight_acc) | torch.isinf(weight_acc), - torch.zeros_like(weight_acc), - weight_acc, - ) - - weight_chunk = torch.exp(chunk_lse_for_merge - lse_new) - weight_chunk = torch.where( - torch.isnan(weight_chunk) | torch.isinf(weight_chunk) | is_chunk_lonely, - torch.zeros_like(weight_chunk), - weight_chunk, - ) - - acc = ( - weight_acc.unsqueeze(-1) * acc - + weight_chunk.unsqueeze(-1) * chunk_output.float() - ) - - lse_acc = lse_new - - output = acc - lse = lse_acc - - is_lonely_final = lse == float("-inf") - - lse = torch.where(is_lonely_final, torch.full_like(lse, float("+inf")), lse) - - if attn_sink is not None: - attn_sink_expanded = attn_sink.view(1, h_q) - exp_diff = torch.exp(attn_sink_expanded - lse) - exp_diff = torch.where( - is_lonely_final, torch.full_like(exp_diff, float("inf")), exp_diff - ) - scale = 1.0 / (1.0 + exp_diff) - output = output * scale.unsqueeze(-1) - - output = torch.where( - is_lonely_final.unsqueeze(-1), torch.zeros_like(output), output - ) - - return output.to(torch.bfloat16), lse - - -# ============================================================================ -# Helper class and functions for token-range based chunking -# ============================================================================ - - -class SlicedKVScope: - """A sliced view of KV scope for a specific token range.""" - - __slots__ = [ - "blocked_k", - "blocked_k_quantized", - "indices_in_kvcache", - "topk_length", - ] - - def __init__(self, blocked_k, blocked_k_quantized, indices_in_kvcache, topk_length): - self.blocked_k = blocked_k - self.blocked_k_quantized = blocked_k_quantized - self.indices_in_kvcache = indices_in_kvcache - self.topk_length = topk_length - - -def slice_kv_scope_for_tokens(orig_scope, start_t: int, end_t: int, s_q: int): - """Slice a KV scope to only include tokens in range [start_t, end_t).""" - if orig_scope is None: - return None - - orig_indices = orig_scope.indices_in_kvcache.reshape( - -1, orig_scope.indices_in_kvcache.size(-1) - ) - sliced_indices = orig_indices[start_t:end_t] - - sliced_topk_length = None - if orig_scope.topk_length is not None: - batch_start = start_t // s_q - batch_end = (end_t + s_q - 1) // s_q - batch_topk_length = orig_scope.topk_length[batch_start:batch_end] - if s_q > 1: - chunk_tokens = end_t - start_t - expanded = batch_topk_length.unsqueeze(1).expand(-1, s_q).reshape(-1) - offset_in_first_batch = start_t % s_q - sliced_topk_length = expanded[ - offset_in_first_batch : offset_in_first_batch + chunk_tokens - ] - else: - sliced_topk_length = batch_topk_length - - return SlicedKVScope( - blocked_k=orig_scope.blocked_k, - blocked_k_quantized=orig_scope.blocked_k_quantized, - indices_in_kvcache=sliced_indices, - topk_length=sliced_topk_length, - ) - - -def compute_token_ranges( - total_tokens: int, - total_topk: int, - d_qk: int, - max_buffer_bytes: int = 2 * 1024 * 1024 * 1024, -) -> List[Tuple[int, int]]: - """Compute token ranges for processing, chunking if buffer would exceed limit.""" - buffer_size_bytes = total_tokens * total_topk * d_qk * 2 - - if buffer_size_bytes <= max_buffer_bytes: - return [(0, total_tokens)] - - max_tokens_per_chunk = max_buffer_bytes // (total_topk * d_qk * 2) - chunk_size = max(1, max_tokens_per_chunk) - - token_ranges = [] - start_t = 0 - while start_t < total_tokens: - end_t = min(start_t + chunk_size, total_tokens) - token_ranges.append((start_t, end_t)) - start_t = end_t - - return token_ranges - - -# ============================================================================ -# Split-K Attention for Large TopK -# ============================================================================ -def run_splitk_unified_attention( - q_reshaped, - gathered_kv, - invalid_mask, - d_v, - sm_scale, - total_tokens, - h_q, - total_topk, - d_qk, - attn_sink=None, - split_k=4, -): - """Run split-K attention for large topk cases.""" - from .triton_mla_kernels_decode_splitk import run_splitk_attention - - return run_splitk_attention( - q_reshaped, - gathered_kv, - invalid_mask, - d_v, - sm_scale, - total_tokens, - h_q, - total_topk, - d_qk, - attn_sink=attn_sink, - split_k=split_k, - ) diff --git a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_dsv4.py b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_dsv4.py deleted file mode 100644 index 429a6e02f..000000000 --- a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_dsv4.py +++ /dev/null @@ -1,1355 +0,0 @@ -""" -Triton MLA Decode Kernels for DSV4 (d_qk=512). - -This module contains DSV4-specific gather+dequant kernels and the main -sparse attention decode entry point for DSV4. -""" - -import os -from typing import Optional, Tuple - -import torch -import triton -import triton.language as tl - -from .triton_mla_kernels_decode_common import ( - _bucket_total_tokens, - _get_workload_size_category, - compute_token_ranges, - run_chunked_attention_triton, - run_splitk_unified_attention, - run_unified_attention, - slice_kv_scope_for_tokens, -) - -# Enable Triton autotune cache persistence -TRITON_CACHE_DIR = os.path.join(os.path.dirname(__file__), ".triton_cache") -os.makedirs(TRITON_CACHE_DIR, exist_ok=True) -os.environ.setdefault("TRITON_CACHE_DIR", TRITON_CACHE_DIR) - -# Constants for DSV4 layout -DSV4_D_QK = 512 -DSV4_D_NOPE = 448 -DSV4_D_ROPE = 64 -DSV4_TILE_SIZE = 64 -DSV4_NUM_TILES = 7 -DSV4_BYTES_PER_TOKEN_DATA = 576 # 448 nope + 128 rope -DSV4_BYTES_PER_TOKEN_SCALE = 8 # 7 scales + 1 padding - -# Performance tuning thresholds (empirically determined) -# These thresholds balance kernel launch overhead vs. computation efficiency -# -# DSV4_USE_FUSED_THRESHOLD: Use 1D fused kernel below this element count -# Rationale: Single kernel launch reduces overhead for small/medium workloads -# Value 150K determined by benchmarking on typical production workloads -DSV4_USE_FUSED_THRESHOLD = 150000 -# -# DSV4_USE_FIXED_KERNEL_THRESHOLD: Use fixed BLOCK_TK=128 kernel below this -# Rationale: Avoids autotune overhead for small workloads where fixed config -# performs well. Value 32K balances autotune benefit vs. overhead -DSV4_USE_FIXED_KERNEL_THRESHOLD = 32768 - - -# ============================================================================ -# DSV4 Gather+Dequant Kernels - Optimized with Batched Scale Loading -# ============================================================================ - - -@triton.autotune( - configs=[ - # This is a pure memory-copy + FP8→BF16 dequant kernel. - # - BLOCK_TK controls how many (token×topk) pairs per block. - # - Larger BLOCK_TK amortizes launch overhead but needs more warps. - # - BLOCK_TK=128 is already validated as the fixed config for small workloads - # (below DSV4_USE_FIXED_KERNEL_THRESHOLD = 32K elements). - # - BLOCK_TK=64/128: good for small/medium workloads (fewer warps, less overhead). - # - BLOCK_TK=256: better bandwidth utilization for large workloads. - triton.Config({"BLOCK_TK": 64}, num_warps=4, num_stages=1), - triton.Config({"BLOCK_TK": 128}, num_warps=4, num_stages=1), - triton.Config({"BLOCK_TK": 256}, num_warps=8, num_stages=1), - ], - key=["total_tokens_bucket", "topk", "workload_size_cat"], -) -@triton.jit -def _gather_dequant_dsv4_kernel( - KV_Cache, - Indices, - TopkLength, - OutputKV, - OutputMask, - total_tokens, - total_tokens_bucket, - topk, - num_blocks, - block_size, - workload_size_cat, - k_offset, - s_q, - stride_kv_block, - stride_idx_t, - stride_idx_k, - stride_out_t, - stride_out_k, - stride_out_d, - stride_mask_t, - stride_mask_k, - BLOCK_TK: tl.constexpr, - D_NOPE: tl.constexpr, - D_ROPE: tl.constexpr, - BYTES_PER_TOKEN_DATA: tl.constexpr, - BYTES_PER_TOKEN_SCALE: tl.constexpr, - TILE_SIZE: tl.constexpr, - HAS_TOPK_LENGTH: tl.constexpr, -): - """Optimized gather + dequant kernel with batched scale loading.""" - pid = tl.program_id(0) - num_tk = total_tokens * topk - - offs_tk = pid * BLOCK_TK + tl.arange(0, BLOCK_TK) - mask_tk = offs_tk < num_tk - - t_idx = offs_tk // topk - k_idx = offs_tk % topk - - idx_ptrs = Indices + t_idx * stride_idx_t + k_idx * stride_idx_k - indices = tl.load(idx_ptrs, mask=mask_tk, other=-1) - - is_invalid = indices == -1 - - if HAS_TOPK_LENGTH: - batch_idx = t_idx // s_q - topk_len = tl.load(TopkLength + batch_idx, mask=mask_tk, other=topk) - is_invalid = is_invalid | (k_idx >= topk_len) - - mask_out_ptrs = ( - OutputMask + t_idx * stride_mask_t + (k_idx + k_offset) * stride_mask_k - ) - tl.store(mask_out_ptrs, is_invalid, mask=mask_tk) - - valid_mask = mask_tk & ~is_invalid - indices_clamped = tl.maximum(indices, 0) - - block_idx = indices_clamped // block_size - offset_in_block = indices_clamped % block_size - - block_idx_64 = block_idx.to(tl.int64) - offset_in_block_64 = offset_in_block.to(tl.int64) - - kv_block_base = KV_Cache + block_idx_64 * stride_kv_block - - nope_rope_offset = offset_in_block_64 * BYTES_PER_TOKEN_DATA - scale_base_offset = ( - block_size * BYTES_PER_TOKEN_DATA + offset_in_block_64 * BYTES_PER_TOKEN_SCALE - ) - - t_idx_64 = t_idx.to(tl.int64) - k_idx_64 = k_idx.to(tl.int64) - stride_out_t_64 = tl.cast(stride_out_t, tl.int64) - stride_out_k_64 = tl.cast(stride_out_k, tl.int64) - out_base_ptrs = ( - OutputKV + t_idx_64 * stride_out_t_64 + (k_idx_64 + k_offset) * stride_out_k_64 - ) - - # Load all 7 scales at once - each scale is at scale_base_offset + tile_idx - scale_ptrs_0 = kv_block_base + scale_base_offset - scale_ptrs_1 = kv_block_base + scale_base_offset + 1 - scale_ptrs_2 = kv_block_base + scale_base_offset + 2 - scale_ptrs_3 = kv_block_base + scale_base_offset + 3 - scale_ptrs_4 = kv_block_base + scale_base_offset + 4 - scale_ptrs_5 = kv_block_base + scale_base_offset + 5 - scale_ptrs_6 = kv_block_base + scale_base_offset + 6 - - scale_uint8_0 = tl.load(scale_ptrs_0, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_1 = tl.load(scale_ptrs_1, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_2 = tl.load(scale_ptrs_2, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_3 = tl.load(scale_ptrs_3, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_4 = tl.load(scale_ptrs_4, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_5 = tl.load(scale_ptrs_5, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_6 = tl.load(scale_ptrs_6, mask=valid_mask, other=127).to(tl.uint8) - - # Convert all scales to bf16 and pre-compute 2D versions - scale_bf16_0 = tl.math.exp2(scale_uint8_0.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_1 = tl.math.exp2(scale_uint8_1.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_2 = tl.math.exp2(scale_uint8_2.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_3 = tl.math.exp2(scale_uint8_3.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_4 = tl.math.exp2(scale_uint8_4.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_5 = tl.math.exp2(scale_uint8_5.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_6 = tl.math.exp2(scale_uint8_6.to(tl.float32) - 127.0).to(tl.bfloat16) - # Pre-compute 2D versions for tile processing - scale_2d_0 = scale_bf16_0[:, None] - scale_2d_1 = scale_bf16_1[:, None] - scale_2d_2 = scale_bf16_2[:, None] - scale_2d_3 = scale_bf16_3[:, None] - scale_2d_4 = scale_bf16_4[:, None] - scale_2d_5 = scale_bf16_5[:, None] - scale_2d_6 = scale_bf16_6[:, None] - - offs_d = tl.arange(0, TILE_SIZE) - - # Pre-compute base pointers for optimization - tile_base = kv_block_base[:, None] + nope_rope_offset[:, None] - out_base = out_base_ptrs[:, None] - valid_mask_2d = valid_mask[:, None] - is_invalid_2d = is_invalid[:, None] - mask_tk_2d = mask_tk[:, None] - - # Process tile 0 - nope_ptrs = tile_base + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_0 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + offs_d[None, :] * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 1 - tile_start_1 = TILE_SIZE - nope_ptrs = tile_base + tile_start_1 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_1 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_1 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 2 - tile_start_2 = 2 * TILE_SIZE - nope_ptrs = tile_base + tile_start_2 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_2 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_2 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 3 - tile_start_3 = 3 * TILE_SIZE - nope_ptrs = tile_base + tile_start_3 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_3 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_3 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 4 - tile_start_4 = 4 * TILE_SIZE - nope_ptrs = tile_base + tile_start_4 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_4 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_4 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 5 - tile_start_5 = 5 * TILE_SIZE - nope_ptrs = tile_base + tile_start_5 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_5 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_5 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 6 - tile_start_6 = 6 * TILE_SIZE - nope_ptrs = tile_base + tile_start_6 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_6 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_6 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process rope - offs_rope = tl.arange(0, D_ROPE) - rope_byte_start = D_NOPE - - rope_lo_ptrs = tile_base + rope_byte_start + offs_rope[None, :] * 2 - rope_hi_ptrs = tile_base + rope_byte_start + offs_rope[None, :] * 2 + 1 - - rope_lo = tl.load(rope_lo_ptrs, mask=valid_mask_2d, other=0).to(tl.uint16) - rope_hi = tl.load(rope_hi_ptrs, mask=valid_mask_2d, other=0).to(tl.uint16) - - rope_uint16 = rope_lo | (rope_hi << 8) - rope_bf16 = rope_uint16.to(tl.bfloat16, bitcast=True) - rope_bf16 = tl.where(is_invalid_2d, 0.0, rope_bf16) - - out_ptrs = out_base + (D_NOPE + offs_rope[None, :]) * stride_out_d - tl.store(out_ptrs, rope_bf16, mask=mask_tk_2d) - - -@triton.jit -def _gather_dequant_dsv4_kernel_fixed_128( - KV_Cache, - Indices, - TopkLength, - OutputKV, - OutputMask, - total_tokens, - total_tokens_bucket, - topk, - num_blocks, - block_size, - k_offset, - s_q, - stride_kv_block, - stride_idx_t, - stride_idx_k, - stride_out_t, - stride_out_k, - stride_out_d, - stride_mask_t, - stride_mask_k, - D_NOPE: tl.constexpr, - D_ROPE: tl.constexpr, - BYTES_PER_TOKEN_DATA: tl.constexpr, - BYTES_PER_TOKEN_SCALE: tl.constexpr, - TILE_SIZE: tl.constexpr, - HAS_TOPK_LENGTH: tl.constexpr, -): - """Fixed-config gather kernel with BLOCK_TK=128 and batched scale loading.""" - BLOCK_TK: tl.constexpr = 128 - pid = tl.program_id(0) - num_tk = total_tokens * topk - - offs_tk = pid * BLOCK_TK + tl.arange(0, BLOCK_TK) - mask_tk = offs_tk < num_tk - - t_idx = offs_tk // topk - k_idx = offs_tk % topk - - idx_ptrs = Indices + t_idx * stride_idx_t + k_idx * stride_idx_k - indices = tl.load(idx_ptrs, mask=mask_tk, other=-1) - - is_invalid = indices == -1 - - if HAS_TOPK_LENGTH: - batch_idx = t_idx // s_q - topk_len = tl.load(TopkLength + batch_idx, mask=mask_tk, other=topk) - is_invalid = is_invalid | (k_idx >= topk_len) - - mask_out_ptrs = ( - OutputMask + t_idx * stride_mask_t + (k_idx + k_offset) * stride_mask_k - ) - tl.store(mask_out_ptrs, is_invalid, mask=mask_tk) - - valid_mask = mask_tk & ~is_invalid - indices_clamped = tl.maximum(indices, 0) - - block_idx = indices_clamped // block_size - offset_in_block = indices_clamped % block_size - - block_idx_64 = block_idx.to(tl.int64) - offset_in_block_64 = offset_in_block.to(tl.int64) - - kv_block_base = KV_Cache + block_idx_64 * stride_kv_block - - nope_rope_offset = offset_in_block_64 * BYTES_PER_TOKEN_DATA - scale_base_offset = ( - block_size * BYTES_PER_TOKEN_DATA + offset_in_block_64 * BYTES_PER_TOKEN_SCALE - ) - - t_idx_64 = t_idx.to(tl.int64) - k_idx_64 = k_idx.to(tl.int64) - stride_out_t_64 = tl.cast(stride_out_t, tl.int64) - stride_out_k_64 = tl.cast(stride_out_k, tl.int64) - out_base_ptrs = ( - OutputKV + t_idx_64 * stride_out_t_64 + (k_idx_64 + k_offset) * stride_out_k_64 - ) - - # Load all 7 scales at once - scale_ptrs_0 = kv_block_base + scale_base_offset - scale_ptrs_1 = kv_block_base + scale_base_offset + 1 - scale_ptrs_2 = kv_block_base + scale_base_offset + 2 - scale_ptrs_3 = kv_block_base + scale_base_offset + 3 - scale_ptrs_4 = kv_block_base + scale_base_offset + 4 - scale_ptrs_5 = kv_block_base + scale_base_offset + 5 - scale_ptrs_6 = kv_block_base + scale_base_offset + 6 - - scale_uint8_0 = tl.load(scale_ptrs_0, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_1 = tl.load(scale_ptrs_1, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_2 = tl.load(scale_ptrs_2, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_3 = tl.load(scale_ptrs_3, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_4 = tl.load(scale_ptrs_4, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_5 = tl.load(scale_ptrs_5, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_6 = tl.load(scale_ptrs_6, mask=valid_mask, other=127).to(tl.uint8) - - # Convert all scales to bf16 and pre-compute 2D versions - scale_bf16_0 = tl.math.exp2(scale_uint8_0.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_1 = tl.math.exp2(scale_uint8_1.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_2 = tl.math.exp2(scale_uint8_2.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_3 = tl.math.exp2(scale_uint8_3.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_4 = tl.math.exp2(scale_uint8_4.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_5 = tl.math.exp2(scale_uint8_5.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_6 = tl.math.exp2(scale_uint8_6.to(tl.float32) - 127.0).to(tl.bfloat16) - # Pre-compute 2D versions for tile processing - scale_2d_0 = scale_bf16_0[:, None] - scale_2d_1 = scale_bf16_1[:, None] - scale_2d_2 = scale_bf16_2[:, None] - scale_2d_3 = scale_bf16_3[:, None] - scale_2d_4 = scale_bf16_4[:, None] - scale_2d_5 = scale_bf16_5[:, None] - scale_2d_6 = scale_bf16_6[:, None] - - offs_d = tl.arange(0, TILE_SIZE) - - # Pre-compute base pointers for optimization - tile_base = kv_block_base[:, None] + nope_rope_offset[:, None] - out_base = out_base_ptrs[:, None] - valid_mask_2d = valid_mask[:, None] - is_invalid_2d = is_invalid[:, None] - mask_tk_2d = mask_tk[:, None] - - # Process tile 0 - nope_ptrs = tile_base + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_0 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + offs_d[None, :] * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 1 - tile_start_1 = TILE_SIZE - nope_ptrs = tile_base + tile_start_1 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_1 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_1 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 2 - tile_start_2 = 2 * TILE_SIZE - nope_ptrs = tile_base + tile_start_2 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_2 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_2 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 3 - tile_start_3 = 3 * TILE_SIZE - nope_ptrs = tile_base + tile_start_3 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_3 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_3 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 4 - tile_start_4 = 4 * TILE_SIZE - nope_ptrs = tile_base + tile_start_4 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_4 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_4 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 5 - tile_start_5 = 5 * TILE_SIZE - nope_ptrs = tile_base + tile_start_5 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_5 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_5 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 6 - tile_start_6 = 6 * TILE_SIZE - nope_ptrs = tile_base + tile_start_6 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_6 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_6 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process rope - offs_rope = tl.arange(0, D_ROPE) - rope_byte_start = D_NOPE - - rope_lo_ptrs = tile_base + rope_byte_start + offs_rope[None, :] * 2 - rope_hi_ptrs = tile_base + rope_byte_start + offs_rope[None, :] * 2 + 1 - - rope_lo = tl.load(rope_lo_ptrs, mask=valid_mask_2d, other=0).to(tl.uint16) - rope_hi = tl.load(rope_hi_ptrs, mask=valid_mask_2d, other=0).to(tl.uint16) - - rope_uint16 = rope_lo | (rope_hi << 8) - rope_bf16 = rope_uint16.to(tl.bfloat16, bitcast=True) - rope_bf16 = tl.where(is_invalid_2d, 0.0, rope_bf16) - - out_ptrs = out_base + (D_NOPE + offs_rope[None, :]) * stride_out_d - tl.store(out_ptrs, rope_bf16, mask=mask_tk_2d) - - -# ============================================================================ -# DSV4 Wrapper Functions -# ============================================================================ - - -def gather_dequant_fp8_dsv4( - kv_cache_quantized: torch.Tensor, - indices: torch.Tensor, - block_size: int, - output_kv: torch.Tensor, - output_mask: torch.Tensor, - k_offset: int = 0, - topk_length: Optional[torch.Tensor] = None, - s_q: int = 1, -) -> bool: - """Unified DSV4 gather+dequant with optional topk_length mask.""" - total_tokens, topk = indices.shape - num_blocks = kv_cache_quantized.shape[0] - - kv_uint8 = kv_cache_quantized.view(torch.uint8) - bytes_per_block = kv_uint8.shape[1] * kv_uint8.shape[2] * kv_uint8.shape[3] - kv_flat = kv_uint8.reshape(num_blocks, bytes_per_block) - - stride_kv_block = kv_uint8.stride(0) - workload_size_cat = _get_workload_size_category(total_tokens, topk) - - grid = lambda meta: (triton.cdiv(total_tokens * topk, meta["BLOCK_TK"]),) - - topk_length_tensor = topk_length if topk_length is not None else output_mask[:1, 0] - has_topk_length = topk_length is not None - - _gather_dequant_dsv4_kernel[grid]( - kv_flat, - indices, - topk_length_tensor, - output_kv, - output_mask, - total_tokens, - _bucket_total_tokens(total_tokens), - topk, - num_blocks, - block_size, - workload_size_cat, - k_offset, - s_q, - stride_kv_block, - indices.stride(0), - indices.stride(1), - output_kv.stride(0), - output_kv.stride(1), - output_kv.stride(2), - output_mask.stride(0), - output_mask.stride(1), - D_NOPE=DSV4_D_NOPE, - D_ROPE=DSV4_D_ROPE, - BYTES_PER_TOKEN_DATA=DSV4_BYTES_PER_TOKEN_DATA, - BYTES_PER_TOKEN_SCALE=DSV4_BYTES_PER_TOKEN_SCALE, - TILE_SIZE=DSV4_TILE_SIZE, - HAS_TOPK_LENGTH=has_topk_length, - ) - return True - - -# ============================================================================ -# DSV4 1D Grid Fused Gather+Dequant Kernel (Optimized - No Empty Blocks) -# Single kernel launch with 1D grid: (num_main_pids + num_extra_pids,) -# ============================================================================ - - -@triton.jit -def _gather_dequant_dsv4_1d_fused_kernel( - # Main KV cache - KV_Cache_Main, - Indices_Main, - TopkLength_Main, - # Extra KV cache - KV_Cache_Extra, - Indices_Extra, - TopkLength_Extra, - # Output - OutputKV, - OutputMask, - # Dimensions - total_tokens, - topk_main, - topk_extra, - num_blocks_main, - num_blocks_extra, - block_size_main, - block_size_extra, - s_q, - # Strides for main - stride_kv_block_main, - stride_idx_t_main, - stride_idx_k_main, - # Strides for extra - stride_kv_block_extra, - stride_idx_t_extra, - stride_idx_k_extra, - # Output strides - stride_out_t, - stride_out_k, - stride_out_d, - stride_mask_t, - stride_mask_k, - # Grid info - num_main_pids, - # Constexpr - BLOCK_TK: tl.constexpr, - D_NOPE: tl.constexpr, - D_ROPE: tl.constexpr, - BYTES_PER_TOKEN_DATA: tl.constexpr, - BYTES_PER_TOKEN_SCALE: tl.constexpr, - TILE_SIZE: tl.constexpr, - HAS_TOPK_LENGTH_MAIN: tl.constexpr, - HAS_TOPK_LENGTH_EXTRA: tl.constexpr, -): - """1D fused gather kernel - single launch, no empty blocks. - - Grid: (num_main_pids + num_extra_pids,) - - pid < num_main_pids: process main cache - - pid >= num_main_pids: process extra cache - - This eliminates empty blocks when main/extra topk differ significantly. - """ - pid = tl.program_id(0) - - # Determine if this is main or extra processing - is_main_pid = pid < num_main_pids - - # Select parameters based on pid - if is_main_pid: - local_pid = pid - topk = topk_main - k_offset = 0 - num_tk = total_tokens * topk_main - KV_Cache = KV_Cache_Main - Indices = Indices_Main - TopkLength = TopkLength_Main - block_size = block_size_main - stride_kv_block = stride_kv_block_main - stride_idx_t = stride_idx_t_main - stride_idx_k = stride_idx_k_main - else: - local_pid = pid - num_main_pids - topk = topk_extra - k_offset = topk_main - num_tk = total_tokens * topk_extra - KV_Cache = KV_Cache_Extra - Indices = Indices_Extra - TopkLength = TopkLength_Extra - block_size = block_size_extra - stride_kv_block = stride_kv_block_extra - stride_idx_t = stride_idx_t_extra - stride_idx_k = stride_idx_k_extra - - # Compute element indices for this block - offs_tk = local_pid * BLOCK_TK + tl.arange(0, BLOCK_TK) - mask_tk = offs_tk < num_tk - - t_idx = offs_tk // topk - k_idx = offs_tk % topk - - # Load indices - idx_ptrs = Indices + t_idx * stride_idx_t + k_idx * stride_idx_k - indices = tl.load(idx_ptrs, mask=mask_tk, other=-1) - - is_invalid = indices == -1 - - # Handle topk_length - need to handle both cases - batch_idx = t_idx // s_q - if is_main_pid: - if HAS_TOPK_LENGTH_MAIN: - topk_len = tl.load(TopkLength + batch_idx, mask=mask_tk, other=topk) - is_invalid = is_invalid | (k_idx >= topk_len) - else: - if HAS_TOPK_LENGTH_EXTRA: - topk_len = tl.load(TopkLength + batch_idx, mask=mask_tk, other=topk) - is_invalid = is_invalid | (k_idx >= topk_len) - - # Store mask - mask_out_ptrs = ( - OutputMask + t_idx * stride_mask_t + (k_idx + k_offset) * stride_mask_k - ) - tl.store(mask_out_ptrs, is_invalid, mask=mask_tk) - - valid_mask = mask_tk & ~is_invalid - indices_clamped = tl.maximum(indices, 0) - - block_idx = indices_clamped // block_size - offset_in_block = indices_clamped % block_size - - block_idx_64 = block_idx.to(tl.int64) - offset_in_block_64 = offset_in_block.to(tl.int64) - - kv_block_base = KV_Cache + block_idx_64 * stride_kv_block - - nope_rope_offset = offset_in_block_64 * BYTES_PER_TOKEN_DATA - scale_base_offset = ( - block_size * BYTES_PER_TOKEN_DATA + offset_in_block_64 * BYTES_PER_TOKEN_SCALE - ) - - t_idx_64 = t_idx.to(tl.int64) - k_idx_64 = k_idx.to(tl.int64) - stride_out_t_64 = tl.cast(stride_out_t, tl.int64) - stride_out_k_64 = tl.cast(stride_out_k, tl.int64) - out_base_ptrs = ( - OutputKV + t_idx_64 * stride_out_t_64 + (k_idx_64 + k_offset) * stride_out_k_64 - ) - - # Load all 7 scales - scale_ptrs_0 = kv_block_base + scale_base_offset - scale_uint8_0 = tl.load(scale_ptrs_0, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_1 = tl.load(scale_ptrs_0 + 1, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_2 = tl.load(scale_ptrs_0 + 2, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_3 = tl.load(scale_ptrs_0 + 3, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_4 = tl.load(scale_ptrs_0 + 4, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_5 = tl.load(scale_ptrs_0 + 5, mask=valid_mask, other=127).to(tl.uint8) - scale_uint8_6 = tl.load(scale_ptrs_0 + 6, mask=valid_mask, other=127).to(tl.uint8) - - scale_bf16_0 = tl.math.exp2(scale_uint8_0.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_1 = tl.math.exp2(scale_uint8_1.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_2 = tl.math.exp2(scale_uint8_2.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_3 = tl.math.exp2(scale_uint8_3.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_4 = tl.math.exp2(scale_uint8_4.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_5 = tl.math.exp2(scale_uint8_5.to(tl.float32) - 127.0).to(tl.bfloat16) - scale_bf16_6 = tl.math.exp2(scale_uint8_6.to(tl.float32) - 127.0).to(tl.bfloat16) - # Pre-compute 2D versions for tile processing - scale_2d_0 = scale_bf16_0[:, None] - scale_2d_1 = scale_bf16_1[:, None] - scale_2d_2 = scale_bf16_2[:, None] - scale_2d_3 = scale_bf16_3[:, None] - scale_2d_4 = scale_bf16_4[:, None] - scale_2d_5 = scale_bf16_5[:, None] - scale_2d_6 = scale_bf16_6[:, None] - - offs_d = tl.arange(0, TILE_SIZE) - - # Pre-compute base pointers for optimization - tile_base = kv_block_base[:, None] + nope_rope_offset[:, None] - out_base = out_base_ptrs[:, None] - valid_mask_2d = valid_mask[:, None] - is_invalid_2d = is_invalid[:, None] - mask_tk_2d = mask_tk[:, None] - - # Process tile 0 - nope_ptrs = tile_base + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_0 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + offs_d[None, :] * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 1 - tile_start_1 = TILE_SIZE - nope_ptrs = tile_base + tile_start_1 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_1 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_1 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 2 - tile_start_2 = 2 * TILE_SIZE - nope_ptrs = tile_base + tile_start_2 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_2 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_2 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 3 - tile_start_3 = 3 * TILE_SIZE - nope_ptrs = tile_base + tile_start_3 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_3 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_3 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 4 - tile_start_4 = 4 * TILE_SIZE - nope_ptrs = tile_base + tile_start_4 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_4 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_4 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 5 - tile_start_5 = 5 * TILE_SIZE - nope_ptrs = tile_base + tile_start_5 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_5 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_5 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process tile 6 - tile_start_6 = 6 * TILE_SIZE - nope_ptrs = tile_base + tile_start_6 + offs_d[None, :] - nope_uint8 = tl.load(nope_ptrs, mask=valid_mask_2d, other=0) - nope_fp8 = nope_uint8.to(tl.float8e4nv, bitcast=True) - nope_bf16 = nope_fp8.to(tl.bfloat16) - dequant = nope_bf16 * scale_2d_6 - dequant = tl.where(is_invalid_2d, 0.0, dequant) - out_ptrs = out_base + (tile_start_6 + offs_d[None, :]) * stride_out_d - tl.store(out_ptrs, dequant, mask=mask_tk_2d) - - # Process rope - offs_rope = tl.arange(0, D_ROPE) - rope_byte_start = D_NOPE - rope_lo_ptrs = tile_base + rope_byte_start + offs_rope[None, :] * 2 - rope_hi_ptrs = tile_base + rope_byte_start + offs_rope[None, :] * 2 + 1 - rope_lo = tl.load(rope_lo_ptrs, mask=valid_mask_2d, other=0).to(tl.uint16) - rope_hi = tl.load(rope_hi_ptrs, mask=valid_mask_2d, other=0).to(tl.uint16) - rope_uint16 = rope_lo | (rope_hi << 8) - rope_bf16 = rope_uint16.to(tl.bfloat16, bitcast=True) - rope_bf16 = tl.where(is_invalid_2d, 0.0, rope_bf16) - out_ptrs = out_base + (D_NOPE + offs_rope[None, :]) * stride_out_d - tl.store(out_ptrs, rope_bf16, mask=mask_tk_2d) - - -def _prepare_kv_cache_flat(kv_cache): - """Helper to prepare KV cache for gather operations. - - Returns: (kv_flat, num_blocks, stride_kv_block) - """ - kv_uint8 = kv_cache.view(torch.uint8) - num_blocks = kv_cache.shape[0] - bytes_per_block = kv_uint8.shape[1] * kv_uint8.shape[2] * kv_uint8.shape[3] - kv_flat = kv_uint8.reshape(num_blocks, bytes_per_block) - stride_kv_block = kv_uint8.stride(0) - return kv_flat, num_blocks, stride_kv_block - - -def _launch_gather_dequant_one_dsv4( - kv_flat, - indices, - topk_length_tensor, - output_kv, - output_mask, - total_tokens, - topk, - num_blocks, - block_size, - k_offset, - s_q, - stride_kv_block, - stride_idx_t, - stride_idx_k, - stride_out_t, - stride_out_k, - stride_out_d, - stride_mask_t, - stride_mask_k, - has_topk_length, -): - """Helper to launch gather+dequant kernel for one KV cache (main or extra). - - This eliminates code duplication between main and extra kernel launches - in the two-kernel path of fused_gather_dequant_fp8_dsv4. - """ - total_elements = total_tokens * topk - - if total_elements < DSV4_USE_FIXED_KERNEL_THRESHOLD: - grid = (triton.cdiv(total_elements, 128),) - _gather_dequant_dsv4_kernel_fixed_128[grid]( - kv_flat, - indices, - topk_length_tensor, - output_kv, - output_mask, - total_tokens, - _bucket_total_tokens(total_tokens), - topk, - num_blocks, - block_size, - k_offset, - s_q, - stride_kv_block, - stride_idx_t, - stride_idx_k, - stride_out_t, - stride_out_k, - stride_out_d, - stride_mask_t, - stride_mask_k, - D_NOPE=DSV4_D_NOPE, - D_ROPE=DSV4_D_ROPE, - BYTES_PER_TOKEN_DATA=DSV4_BYTES_PER_TOKEN_DATA, - BYTES_PER_TOKEN_SCALE=DSV4_BYTES_PER_TOKEN_SCALE, - TILE_SIZE=DSV4_TILE_SIZE, - HAS_TOPK_LENGTH=has_topk_length, - num_warps=8, - num_stages=2, - ) - else: - workload_cat = _get_workload_size_category(total_tokens, topk) - grid = lambda meta: (triton.cdiv(total_elements, meta["BLOCK_TK"]),) - _gather_dequant_dsv4_kernel[grid]( - kv_flat, - indices, - topk_length_tensor, - output_kv, - output_mask, - total_tokens, - _bucket_total_tokens(total_tokens), - topk, - num_blocks, - block_size, - workload_cat, - k_offset, - s_q, - stride_kv_block, - stride_idx_t, - stride_idx_k, - stride_out_t, - stride_out_k, - stride_out_d, - stride_mask_t, - stride_mask_k, - D_NOPE=DSV4_D_NOPE, - D_ROPE=DSV4_D_ROPE, - BYTES_PER_TOKEN_DATA=DSV4_BYTES_PER_TOKEN_DATA, - BYTES_PER_TOKEN_SCALE=DSV4_BYTES_PER_TOKEN_SCALE, - TILE_SIZE=DSV4_TILE_SIZE, - HAS_TOPK_LENGTH=has_topk_length, - ) - - -def truly_fused_gather_dequant_fp8_dsv4( - kv_cache_main, - indices_main, - block_size_main, - topk_length_main, - kv_cache_extra, - indices_extra, - block_size_extra, - topk_length_extra, - output_kv, - output_mask, - s_q=1, -): - """Truly fused DSV4 gather - single kernel launch with 1D grid (no empty blocks).""" - total_tokens, topk_main = indices_main.shape - topk_extra = indices_extra.shape[1] - b = total_tokens // s_q # batch size - - kv_flat_main, num_blocks_main, stride_kv_block_main = _prepare_kv_cache_flat( - kv_cache_main - ) - kv_flat_extra, num_blocks_extra, stride_kv_block_extra = _prepare_kv_cache_flat( - kv_cache_extra - ) - - has_topk_length_main = topk_length_main is not None - has_topk_length_extra = topk_length_extra is not None - - # Always use int32 tensors for topk_length to avoid type mismatch in Triton - if has_topk_length_main: - topk_length_main_tensor = topk_length_main - else: - topk_length_main_tensor = torch.full( - (b,), topk_main, dtype=torch.int32, device=indices_main.device - ) - - if has_topk_length_extra: - topk_length_extra_tensor = topk_length_extra - else: - topk_length_extra_tensor = torch.full( - (b,), topk_extra, dtype=torch.int32, device=indices_extra.device - ) - - stride_idx_t_main, stride_idx_k_main = indices_main.stride(0), indices_main.stride( - 1 - ) - stride_idx_t_extra, stride_idx_k_extra = indices_extra.stride( - 0 - ), indices_extra.stride(1) - stride_out_t, stride_out_k, stride_out_d = ( - output_kv.stride(0), - output_kv.stride(1), - output_kv.stride(2), - ) - stride_mask_t, stride_mask_k = output_mask.stride(0), output_mask.stride(1) - - BLOCK_TK = 128 - - # Calculate grid sizes - 1D grid with exact number of needed blocks - num_elements_main = total_tokens * topk_main - num_elements_extra = total_tokens * topk_extra - num_main_pids = triton.cdiv(num_elements_main, BLOCK_TK) - num_extra_pids = triton.cdiv(num_elements_extra, BLOCK_TK) - - # 1D grid: (num_main_pids + num_extra_pids,) - no empty blocks! - grid = (num_main_pids + num_extra_pids,) - - _gather_dequant_dsv4_1d_fused_kernel[grid]( - kv_flat_main, - indices_main, - topk_length_main_tensor, - kv_flat_extra, - indices_extra, - topk_length_extra_tensor, - output_kv, - output_mask, - total_tokens, - topk_main, - topk_extra, - num_blocks_main, - num_blocks_extra, - block_size_main, - block_size_extra, - s_q, - stride_kv_block_main, - stride_idx_t_main, - stride_idx_k_main, - stride_kv_block_extra, - stride_idx_t_extra, - stride_idx_k_extra, - stride_out_t, - stride_out_k, - stride_out_d, - stride_mask_t, - stride_mask_k, - num_main_pids, - BLOCK_TK=BLOCK_TK, - D_NOPE=DSV4_D_NOPE, - D_ROPE=DSV4_D_ROPE, - BYTES_PER_TOKEN_DATA=DSV4_BYTES_PER_TOKEN_DATA, - BYTES_PER_TOKEN_SCALE=DSV4_BYTES_PER_TOKEN_SCALE, - TILE_SIZE=DSV4_TILE_SIZE, - HAS_TOPK_LENGTH_MAIN=has_topk_length_main, - HAS_TOPK_LENGTH_EXTRA=has_topk_length_extra, - num_warps=8, - num_stages=2, - ) - return True - - -def fused_gather_dequant_fp8_dsv4( - kv_cache_main, - indices_main, - block_size_main, - topk_length_main, - kv_cache_extra, - indices_extra, - block_size_extra, - topk_length_extra, - output_kv, - output_mask, - s_q=1, -): - """Fused DSV4 gather - uses 1D fused kernel for small workloads, two kernels for large.""" - has_topk_length_main = topk_length_main is not None - has_topk_length_extra = topk_length_extra is not None - - total_tokens, topk_main = indices_main.shape - topk_extra = indices_extra.shape[1] - total_elements = total_tokens * (topk_main + topk_extra) - - # Use fused 2D grid kernel only for small workloads where kernel launch overhead matters - # For large workloads, the two-kernel approach is more efficient - USE_FUSED_THRESHOLD = DSV4_USE_FUSED_THRESHOLD - - # IMPORTANT: Disable fused kernel when topk_length settings differ between main and extra - # The 1D fused kernel has issues with runtime conditional handling when - # HAS_TOPK_LENGTH_MAIN != HAS_TOPK_LENGTH_EXTRA, causing incorrect results in extra part. - # Only use fused kernel when both have same topk_length setting. - topk_length_settings_match = has_topk_length_main == has_topk_length_extra - use_fused = total_elements < USE_FUSED_THRESHOLD and topk_length_settings_match - - if use_fused: - return truly_fused_gather_dequant_fp8_dsv4( - kv_cache_main, - indices_main, - block_size_main, - topk_length_main, - kv_cache_extra, - indices_extra, - block_size_extra, - topk_length_extra, - output_kv, - output_mask, - s_q, - ) - - # Use original two-kernel approach for large workloads - kv_flat_main, num_blocks_main, stride_kv_block_main = _prepare_kv_cache_flat( - kv_cache_main - ) - kv_flat_extra, num_blocks_extra, stride_kv_block_extra = _prepare_kv_cache_flat( - kv_cache_extra - ) - - topk_length_main_tensor = ( - topk_length_main if has_topk_length_main else output_mask[:1, 0] - ) - topk_length_extra_tensor = ( - topk_length_extra if has_topk_length_extra else output_mask[:1, 0] - ) - - stride_idx_t_main, stride_idx_k_main = indices_main.stride(0), indices_main.stride( - 1 - ) - stride_idx_t_extra, stride_idx_k_extra = indices_extra.stride( - 0 - ), indices_extra.stride(1) - stride_out_t, stride_out_k, stride_out_d = ( - output_kv.stride(0), - output_kv.stride(1), - output_kv.stride(2), - ) - stride_mask_t, stride_mask_k = output_mask.stride(0), output_mask.stride(1) - - # Launch main kernel - _launch_gather_dequant_one_dsv4( - kv_flat_main, - indices_main, - topk_length_main_tensor, - output_kv, - output_mask, - total_tokens, - topk_main, - num_blocks_main, - block_size_main, - 0, - s_q, - stride_kv_block_main, - stride_idx_t_main, - stride_idx_k_main, - stride_out_t, - stride_out_k, - stride_out_d, - stride_mask_t, - stride_mask_k, - has_topk_length_main, - ) - - # Launch extra kernel - _launch_gather_dequant_one_dsv4( - kv_flat_extra, - indices_extra, - topk_length_extra_tensor, - output_kv, - output_mask, - total_tokens, - topk_extra, - num_blocks_extra, - block_size_extra, - topk_main, - s_q, - stride_kv_block_extra, - stride_idx_t_extra, - stride_idx_k_extra, - stride_out_t, - stride_out_k, - stride_out_d, - stride_mask_t, - stride_mask_k, - has_topk_length_extra, - ) - - return True - - -def triton_sparse_attn_decode_dsv4( - q: torch.Tensor, - kv_scope, - extra_kv_scope, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Sparse attention decode for DSV4 (d_qk=512).""" - assert kv_scope is not None - b, s_q, h_q, d_qk = q.shape - assert d_qk == DSV4_D_QK, f"Expected d_qk={DSV4_D_QK} for DSV4, got {d_qk}" - total_tokens = b * s_q - - topk_main = kv_scope.indices_in_kvcache.size(-1) - topk_extra = ( - extra_kv_scope.indices_in_kvcache.size(-1) if extra_kv_scope is not None else 0 - ) - total_topk = topk_main + topk_extra - - token_ranges = compute_token_ranges(total_tokens, total_topk, d_qk) - - if len(token_ranges) == 1: - return _triton_sparse_attn_decode_dsv4_impl( - q, kv_scope, extra_kv_scope, sm_scale, d_v, attn_sink - ) - - outputs = [] - lses = [] - - for start_t, end_t in token_ranges: - chunk_tokens = end_t - start_t - q_chunk = q.reshape(total_tokens, h_q, d_qk)[start_t:end_t] - q_input = q_chunk.reshape(chunk_tokens, 1, h_q, d_qk) - chunk_kv_scope = slice_kv_scope_for_tokens(kv_scope, start_t, end_t, s_q) - chunk_extra_kv_scope = slice_kv_scope_for_tokens( - extra_kv_scope, start_t, end_t, s_q - ) - - chunk_out, chunk_lse = _triton_sparse_attn_decode_dsv4_impl( - q_input, chunk_kv_scope, chunk_extra_kv_scope, sm_scale, d_v, attn_sink - ) - - outputs.append(chunk_out.reshape(chunk_tokens, h_q, d_v)) - lses.append(chunk_lse.reshape(chunk_tokens, h_q)) - - output = torch.cat(outputs, dim=0).reshape(b, s_q, h_q, d_v) - lse = torch.cat(lses, dim=0).reshape(b, s_q, h_q).transpose(1, 2) - - return output, lse - - -def _triton_sparse_attn_decode_dsv4_impl( - q: torch.Tensor, - kv_scope, - extra_kv_scope, - sm_scale: float, - d_v: int = 512, - attn_sink: Optional[torch.Tensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Internal implementation of sparse attention decode for DSV4. - - Assumes KV cache is always FP8 quantized (blocked_k_quantized is not None). - """ - assert kv_scope is not None - b, s_q, h_q, d_qk = q.shape - total_tokens = b * s_q - - topk_main = kv_scope.indices_in_kvcache.size(-1) - topk_extra = ( - extra_kv_scope.indices_in_kvcache.size(-1) if extra_kv_scope is not None else 0 - ) - total_topk = topk_main + topk_extra - - gathered_kv = torch.empty( - total_tokens, total_topk, d_qk, dtype=torch.bfloat16, device=q.device - ) - invalid_mask = torch.empty( - total_tokens, total_topk, dtype=torch.bool, device=q.device - ) - - block_size_main = kv_scope.blocked_k.shape[1] - indices_main = kv_scope.indices_in_kvcache.reshape(total_tokens, topk_main) - - if extra_kv_scope is not None: - # Fused gather for both main and extra scope - block_size_extra = extra_kv_scope.blocked_k.shape[1] - indices_extra = extra_kv_scope.indices_in_kvcache.reshape( - total_tokens, topk_extra - ) - fused_gather_dequant_fp8_dsv4( - kv_scope.blocked_k_quantized, - indices_main, - block_size_main, - kv_scope.topk_length, - extra_kv_scope.blocked_k_quantized, - indices_extra, - block_size_extra, - extra_kv_scope.topk_length, - gathered_kv, - invalid_mask, - s_q, - ) - else: - # Single gather for main scope only - gather_dequant_fp8_dsv4( - kv_scope.blocked_k_quantized, - indices_main, - block_size_main, - gathered_kv, - invalid_mask, - 0, - kv_scope.topk_length, - s_q, - ) - - q_reshaped = q.to(torch.bfloat16).reshape(total_tokens, h_q, d_qk) - - if not q_reshaped.is_contiguous(): - q_reshaped = q_reshaped.contiguous() - - # Use splitk for large topk to reduce register pressure - if total_topk >= 8192: - # Adaptive split_k selection for optimal performance - # split_k=3 is optimal for topk >= 16384 based on benchmarking - if total_topk >= 16384: - split_k = 3 - else: - split_k = 2 - output, lse = run_splitk_unified_attention( - q_reshaped, - gathered_kv, - invalid_mask, - d_v, - sm_scale, - total_tokens, - h_q, - total_topk, - d_qk, - attn_sink=attn_sink, - split_k=split_k, - ) - elif total_topk <= 65536: - output, lse = run_unified_attention( - q_reshaped, - gathered_kv, - invalid_mask, - d_v, - sm_scale, - total_tokens, - h_q, - total_topk, - d_qk, - attn_sink=attn_sink, - ) - else: - output, lse = run_chunked_attention_triton( - q_reshaped, - gathered_kv, - invalid_mask, - d_v, - sm_scale, - total_tokens, - h_q, - total_topk, - d_qk, - attn_sink=attn_sink, - chunk_size=32768, - ) - - return output.view(b, s_q, h_q, d_v), lse.view(b, s_q, h_q).transpose(1, 2) diff --git a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_fused.py b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_fused.py index 41c29efb1..1cc82a989 100644 --- a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_fused.py +++ b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_fused.py @@ -6,17 +6,9 @@ This module implements a fused kernel that combines: 2. Dequant: FP8 to BF16 dequantization 3. Attention: Compute attention scores and output -Benefits for workloads without extra scope: -- Eliminates intermediate buffer (gathered_kv) write/read -- Reduces kernel launch overhead (1 kernel instead of 2) -- Better cache utilization - Supports: - DSV4 (d_qk=512): 7 tiles of 64, uint8 scales - All configs: with/without topk_length, with/without attn_sink - -OPTIMIZED VERSION: Reduced code duplication in dual-scope kernel by using -a helper function for KV block processing. """ from typing import Optional, Tuple @@ -25,7 +17,16 @@ import torch import triton import triton.language as tl -from .triton_mla_kernels_decode_common import _bucket_total_tokens + +def _bucket_total_tokens(total_tokens: int) -> int: + """Round total_tokens up to the nearest power of 2 for autotune key stability.""" + if total_tokens <= 0: + return 1 + n = 1 + while n < total_tokens: + n <<= 1 + return n + # ============================================================================ # Constants for DSV4 layout @@ -35,15 +36,93 @@ DSV4_D_NOPE = 448 DSV4_D_ROPE = 64 DSV4_D_V = 512 DSV4_TILE_SIZE = 64 + +# ============================================================================ +# Dispatch thresholds for split-K decision +# ============================================================================ +# Dual-scope topk threshold. +# Split-K is more beneficial for larger topk due to more work per token. +DUAL_SCOPE_SPLITK_TOPK_THRESHOLD = 2048 + +# Token thresholds for split-K vs no-splitk decision: +# - Below these thresholds, split-K provides better GPU utilization. +# - Above these thresholds, the combine kernel overhead dominates. +NOSPLITK_TOKEN_THRESHOLD_LOW_TOPK = ( + 64 # For total_topk < DUAL_SCOPE_SPLITK_TOPK_THRESHOLD +) + +# Small batch threshold: below this, split-K=4/8 for parallelism +SMALL_BATCH_TOKEN_THRESHOLD = 8 + +# Topk threshold for split-K value selection within the split-K path +SPLITK_HIGH_TOPK_THRESHOLD = 512 + +# ============================================================================ +# Shared split-K decision logic for dual-scope kernels +# ============================================================================ + + +def _decide_splitk_dual_scope(total_tokens: int, h_q: int, total_topk: int) -> int: + """Decide the split_k value for dual-scope attention. + + Returns: + split_k value (0 means no split-K, use non-splitk kernel). + """ + # Conditions under which split-K is beneficial: + use_splitk_for_small_bs = total_tokens <= SMALL_BATCH_TOKEN_THRESHOLD and ( + h_q >= 128 or total_topk >= 1024 + ) + use_splitk_for_h64_large_topk = ( + h_q <= 64 + and total_topk >= 1024 + and total_tokens > SMALL_BATCH_TOKEN_THRESHOLD + and total_tokens <= 128 + ) + use_splitk_for_large_topk = ( + total_tokens > NOSPLITK_TOKEN_THRESHOLD_LOW_TOPK + and total_topk >= DUAL_SCOPE_SPLITK_TOPK_THRESHOLD + ) + # For large h_q, the non-splitk grid has very few blocks + # in the H dimension, leading to low GPU utilization. + use_splitk_for_large_hq = ( + h_q > 64 and total_tokens > SMALL_BATCH_TOKEN_THRESHOLD and total_topk >= 256 + ) + + if not ( + use_splitk_for_small_bs + or use_splitk_for_h64_large_topk + or use_splitk_for_large_topk + or use_splitk_for_large_hq + ): + return 0 # No split-K + + # Select split_k value based on workload characteristics. + # Higher topk benefits from more splits; lower topk needs fewer to + # avoid combine overhead. + if total_tokens <= SMALL_BATCH_TOKEN_THRESHOLD: + if total_topk >= SPLITK_HIGH_TOPK_THRESHOLD and total_tokens <= 4: + return 8 + return 4 + elif use_splitk_for_large_hq: + if total_topk >= SPLITK_HIGH_TOPK_THRESHOLD: + return 4 + return 2 + elif use_splitk_for_h64_large_topk: + return 2 + else: + return _select_split_k(total_topk, h_q, total_tokens) + + DSV4_NUM_TILES = 7 DSV4_BYTES_PER_TOKEN_DATA = 576 # 448 nope + 128 rope DSV4_BYTES_PER_TOKEN_SCALE = 8 # 7 scales + 1 padding - # ============================================================================ # Helper: Process KV block and compute QK scores + accumulator update # This is the core computation shared by both single and dual scope kernels # ============================================================================ + + @triton.jit def _process_kv_block_aggressive( # KV cache parameters @@ -198,20 +277,20 @@ def _process_kv_block_aggressive( # ============================================================================ # DSV4 Fused Gather+Dequant+Attention Kernel (Single Scope) # ============================================================================ + + @triton.autotune( configs=[ - # Fused gather+dequant+attention kernel. - # Two axes: BLOCK_H × BLOCK_N, with BLOCK_N being the key perf knob - # for h_q=64 where fewer BLOCK_H values affect the grid. - # BLOCK_N=64: better for large topk (less register pressure per iter). - # BLOCK_N=128: better for small topk (fewer iterations). - # num_warps=4: fused kernel is compute-bound. + triton.Config({"BLOCK_H": 16, "BLOCK_N": 32}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 16, "BLOCK_N": 64}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 16, "BLOCK_N": 128}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_H": 64, "BLOCK_N": 32}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 64, "BLOCK_N": 64}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 64, "BLOCK_N": 128}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 128, "BLOCK_N": 64}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 128, "BLOCK_N": 128}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_H": 64, "BLOCK_N": 64}, num_warps=8, num_stages=1), + triton.Config({"BLOCK_H": 128, "BLOCK_N": 64}, num_warps=8, num_stages=1), ], key=["total_tokens_bucket", "h_q", "topk"], ) @@ -494,9 +573,8 @@ def _fused_gather_attn_dsv4_kernel( tl.store(lse_ptrs, lse, mask=mask_h) -# Threshold for disabling AMD buffer_ops optimization -# When KV cache size exceeds INT32_MAX, buffer_ops can cause int32 overflow -# INT32_MAX = 2^31 - 1 = 2,147,483,647 bytes (~2GB) +# Threshold for disabling buffer_ops optimization +# When KV cache size exceeds this threshold, buffer_ops may overflow BUFFER_OPS_DISABLE_THRESHOLD = 2 * 1024 * 1024 * 1024 # 2GB @@ -739,34 +817,26 @@ def fused_gather_attn_decode_dsv4( return output, lse -# Uses helper function to eliminate code duplication # ============================================================================ def _prune_dual_scope_configs(configs, named_args, **kwargs): - """Prune configs where BLOCK_H > h_q for the dual-scope kernel. + """Prune autotune configs for the dual-scope kernel. - When BLOCK_H > h_q, cdiv(h_q, BLOCK_H) = 1 regardless of BLOCK_H value, - so larger BLOCK_H gives the same grid but may have worse register allocation. - Keep only the smallest BLOCK_H that gives cdiv(h_q, BLOCK_H) = 1, plus - any BLOCK_H <= h_q configs. - - For h_q=64: keep BLOCK_H <= 64 (removes BLOCK_H=128 which gives same grid) - For h_q=128: keep all (all give different grid sizes) + For h_q <= 64: restrict to BLOCK_H=16 only (BLOCK_H >= 32 causes + precision issues in online softmax due to different MFMA reduction orders). + For h_q > 64: prune BLOCK_H > h_q (same grid size, worse register usage). """ h_q = named_args.get("h_q", 128) - pruned = [c for c in configs if c.kwargs.get("BLOCK_H", 16) <= h_q] + if h_q <= 64: + pruned = [c for c in configs if c.kwargs.get("BLOCK_H", 16) <= 16] + else: + pruned = [c for c in configs if c.kwargs.get("BLOCK_H", 16) <= h_q] return pruned if pruned else configs @triton.autotune( configs=[ - # Dual-scope fused gather+dequant+attention. - # Three axes: BLOCK_H × BLOCK_N × (warps, stages). - # - BLOCK_H: {16, 32, 64, 128} covers h_q=64 and h_q=128. - # - BLOCK_N: {64, 128}. BLOCK_N=64 better for large topk, 128 for small topk. - # - _prune_dual_scope_configs removes BLOCK_H > h_q configs (e.g. BLOCK_H=128 - # is pruned when h_q=64 since it gives the same grid as BLOCK_H=64). # warps=4: baseline configs triton.Config({"BLOCK_H": 16, "BLOCK_N": 64}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 16, "BLOCK_N": 128}, num_warps=4, num_stages=1), @@ -1173,47 +1243,22 @@ def _fused_gather_attn_dsv4_dual_scope_kernel( tl.store(lse_ptrs, lse, mask=mask_h) -def _prune_splitk_configs(configs, named_args, **kwargs): - """Prune BLOCK_H=16 configs for large batch sizes to avoid CU oversubscription. - - With h_q=128 and BLOCK_H=16, the grid has cdiv(128,16)=8 H-blocks. - At bs=32 with split_k=2, this creates 8*32*2=512 blocks (200% CU), - causing performance regression from oversubscription. - - For small batch sizes (bucket <= 8), BLOCK_H=16 provides better - parallelism and is ~10% faster in CUDA graph replay. - """ - total_tokens_bucket = named_args.get("total_tokens_bucket", 32) - if total_tokens_bucket > 8: - # Remove BLOCK_H=16 configs for large batch sizes - pruned = [c for c in configs if c.kwargs.get("BLOCK_H", 32) > 16] - if pruned: - return pruned - return configs - - # ============================================================================ # Split-K Kernel for Dual Scope # ============================================================================ + + @triton.autotune( configs=[ - # Split-K dual-scope fused kernel. - # - Split-K adds parallelism in K dim (2-8 splits). - # - BLOCK_N={64,128}: BLOCK_N=64 better for large topk_per_split. - # - num_warps=4: compute-bound fused kernel. - # - BLOCK_H={16,64}: covers h_q=64 and h_q=128. + # BLOCK_H=16 only (BLOCK_H >= 32 causes precision issues). + triton.Config({"BLOCK_H": 16, "BLOCK_N": 32}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 16, "BLOCK_N": 64}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 16, "BLOCK_N": 128}, num_warps=4, num_stages=1), - triton.Config({"BLOCK_H": 64, "BLOCK_N": 64}, num_warps=4, num_stages=1), - triton.Config({"BLOCK_H": 64, "BLOCK_N": 128}, num_warps=4, num_stages=1), - triton.Config({"BLOCK_H": 128, "BLOCK_N": 64}, num_warps=4, num_stages=1), - triton.Config({"BLOCK_H": 128, "BLOCK_N": 128}, num_warps=4, num_stages=1), - # BLOCK_H=32: critical for cc=32 with h_q=128 (gives 256 blocks with split_k=2) - triton.Config({"BLOCK_H": 32, "BLOCK_N": 64}, num_warps=4, num_stages=1), - triton.Config({"BLOCK_H": 32, "BLOCK_N": 128}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_H": 16, "BLOCK_N": 64}, num_warps=8, num_stages=1), + triton.Config({"BLOCK_H": 16, "BLOCK_N": 128}, num_warps=8, num_stages=1), + triton.Config({"BLOCK_H": 16, "BLOCK_N": 32}, num_warps=8, num_stages=1), ], key=["total_tokens_bucket", "h_q", "topk_per_split"], - prune_configs_by={"early_config_prune": _prune_splitk_configs}, ) @triton.jit def _fused_gather_attn_dsv4_dual_scope_splitk_kernel( @@ -1654,58 +1699,12 @@ def fused_gather_attn_decode_dsv4_dual_scope( or kv_cache_size_extra > BUFFER_OPS_DISABLE_THRESHOLD ) - # When force_no_splitk is set, skip the split-K decision and fall - # through to the non-splitk kernel path below. - use_splitk = not force_no_splitk - - # Use Split-K for dual scope in these cases: - # 1. Small batch sizes with h_q=128 or large topk to increase GPU parallelism - # 2. Large topk (>= 2048) with medium/large batch sizes - # 3. NEW: h_q=64 + large topk (>=1024) + medium batch sizes (~21% improvement) - SPLITK_DUAL_SCOPE_TOPK_THRESHOLD = 2048 - # For small bs, only use splitk when h_q=128 or total_topk >= 1024 - use_splitk_for_small_bs = total_tokens <= 8 and (h_q >= 128 or total_topk >= 1024) - # NEW: For h_q=64 with large topk, splitk is beneficial for medium batch sizes - # Only for tokens <= 128 based on benchmarking (bs=64 shows 13% improvement) - use_splitk_for_h64_large_topk = ( - h_q <= 64 and total_topk >= 1024 and total_tokens > 8 and total_tokens <= 128 + split_k = ( + 0 + if force_no_splitk + else _decide_splitk_dual_scope(total_tokens, h_q, total_topk) ) - use_splitk_for_large_topk = ( - total_tokens > 64 and total_topk >= SPLITK_DUAL_SCOPE_TOPK_THRESHOLD - ) - # For h_q > 64 (e.g. h_q=128), the non-splitk grid has very few blocks - # in the H dimension, leading to low GPU utilization at medium batch sizes. - use_splitk_for_large_hq = h_q > 64 and total_tokens > 8 and total_topk >= 256 - if use_splitk and ( - use_splitk_for_small_bs - or use_splitk_for_h64_large_topk - or use_splitk_for_large_topk - or use_splitk_for_large_hq - ): - # Select split_k based on workload and total_topk. - # CUDA graph replay benchmarks show optimal split_k depends on both: - # - High topk (>=512, c4 layers): more splits needed to parallelize - # - Low topk (<512, c128 layers): fewer splits, less combine overhead - if total_tokens <= 8: - if total_topk >= 512 and total_tokens <= 4: - # High topk + very small bs: split_k=8 is 8-33% faster than sk=4 - split_k = 8 - else: - # split_k=4 gives 2x more blocks than split_k=2 - split_k = 4 - elif use_splitk_for_large_hq: - # For h_q > 64 with bs > 8: - if total_topk >= 512: - # High topk: split_k=4 for all medium/large bs - split_k = 4 - else: - # Low topk: split_k=2 is sufficient - split_k = 2 - elif use_splitk_for_h64_large_topk: - # For h_q=64 + large topk + medium bs, split_k=2 is optimal - split_k = 2 - else: - split_k = _select_split_k(total_topk, h_q, total_tokens) + if split_k > 0: topk_per_split = (total_topk + split_k - 1) // split_k partial_output = torch.empty( @@ -1926,21 +1925,20 @@ def fused_gather_attn_decode_dsv4_dual_scope( # Split-K Optimization for Large TopK (>= 8192) # ============================================================================ SPLITK_TOPK_THRESHOLD = 8192 -SPLITK_DEFAULT = 4 @triton.autotune( configs=[ - # Split-K fused kernel for large topk (≥8192). - # - BLOCK_N={16,32}: small blocks for scattered FP8 KV access pattern. - # - num_warps=4: balanced for fused dequant+attention compute. - # - BLOCK_H={16,64}: covers h_q=64 and h_q=128. triton.Config({"BLOCK_H": 16, "BLOCK_N": 16}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 16, "BLOCK_N": 32}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_H": 16, "BLOCK_N": 64}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 64, "BLOCK_N": 16}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 64, "BLOCK_N": 32}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_H": 64, "BLOCK_N": 64}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 128, "BLOCK_N": 16}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 128, "BLOCK_N": 32}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_H": 64, "BLOCK_N": 32}, num_warps=8, num_stages=1), + triton.Config({"BLOCK_H": 128, "BLOCK_N": 32}, num_warps=8, num_stages=1), ], key=["total_tokens_bucket", "h_q", "topk_per_split"], ) @@ -2390,11 +2388,7 @@ def _combine_splitk_kernel( @triton.autotune( configs=[ - # Simple reduce kernel (weighted sum of 8 splits). - # - BLOCK_D=512: covers d_v=512 in one pass (no D-dimension loop). - # - num_warps=8: memory-bound reduce benefits from more warps. - # - split_k=8 is only used at very small batch sizes (≤4 tokens), - # so BLOCK_H=16/32/64 covers the relevant parallelism range. + triton.Config({"BLOCK_H": 16, "BLOCK_D": 512}, num_warps=4, num_stages=1), triton.Config({"BLOCK_H": 16, "BLOCK_D": 512}, num_warps=8, num_stages=1), triton.Config({"BLOCK_H": 32, "BLOCK_D": 512}, num_warps=8, num_stages=1), triton.Config({"BLOCK_H": 64, "BLOCK_D": 512}, num_warps=8, num_stages=1), @@ -2731,12 +2725,11 @@ def _select_split_k(topk: int, h_q: int, total_tokens: int = 64) -> int: the topk dimension. Larger split_k increases parallelism but also increases the overhead of the combine kernel. - Updated heuristics based on benchmarking with optimized BLOCK_N configs: - - For large topk (>= 16384): split_k=4 provides good balance with existing combine kernel - - For medium topk (8192-16383): split_k=4 - - For small topk (< 8192): split_k=2 + Heuristics: + - For large topk (>= SPLITK_TOPK_THRESHOLD): split_k=4 + - For small topk (< SPLITK_TOPK_THRESHOLD): split_k=2 """ - if topk >= 8192: + if topk >= SPLITK_TOPK_THRESHOLD: return 4 else: return 2 @@ -2745,12 +2738,13 @@ def _select_split_k(topk: int, h_q: int, total_tokens: int = 64) -> int: # ============================================================================ # Low-overhead buffer pool for splitk operations # ============================================================================ + + class SplitKBufferPool: """ Pre-allocated buffer pool for split-K intermediate tensors. - Caches partial_output and partial_lse buffers to avoid repeated allocations. - Output buffers are always freshly allocated to ensure correctness. + Caches intermediate buffers to avoid repeated allocations. """ _buffers = {} @@ -2809,7 +2803,6 @@ def fused_gather_attn_decode_dsv4_dual_scope_low_overhead( to minimize Python overhead, which is significant for small batch sizes. The kernel computation is identical to the original version. - Output buffers are always freshly allocated to ensure correctness. """ total_tokens, h_q, d_qk = q.shape topk_main = indices_main.shape[1] @@ -2839,25 +2832,9 @@ def fused_gather_attn_decode_dsv4_dual_scope_low_overhead( indices_extra = indices_extra.contiguous() # Determine split_k - SPLITK_DUAL_SCOPE_TOPK_THRESHOLD = 2048 - use_splitk_for_small_bs = total_tokens <= 8 and (h_q >= 128 or total_topk >= 1024) - use_splitk_for_h64_large_topk = ( - h_q <= 64 and total_topk >= 1024 and total_tokens > 8 and total_tokens <= 128 - ) - use_splitk_for_large_topk = ( - total_tokens > 64 and total_topk >= SPLITK_DUAL_SCOPE_TOPK_THRESHOLD - ) - # For h_q > 64 (e.g. h_q=128), the non-splitk grid has very few blocks - # in the H dimension (cdiv(128,64)=2), leading to low GPU utilization - # at medium batch sizes. Split-K doubles the parallelism. - use_splitk_for_large_hq = h_q > 64 and total_tokens > 8 and total_topk >= 256 + split_k = _decide_splitk_dual_scope(total_tokens, h_q, total_topk) - if not ( - use_splitk_for_small_bs - or use_splitk_for_h64_large_topk - or use_splitk_for_large_topk - or use_splitk_for_large_hq - ): + if split_k == 0: # Fall back to non-splitk version return fused_gather_attn_decode_dsv4_dual_scope( q, @@ -2872,32 +2849,9 @@ def fused_gather_attn_decode_dsv4_dual_scope_low_overhead( topk_length_extra, attn_sink, s_q, + force_no_splitk=True, ) - # Select split_k based on workload and total_topk. - # CUDA graph replay benchmarks show optimal split_k depends on both: - # - High topk (>=512, c4 layers): more splits needed to parallelize - # - Low topk (<512, c128 layers): fewer splits, less combine overhead - if total_tokens <= 8: - if total_topk >= 512 and total_tokens <= 4: - # High topk + very small bs: split_k=8 is 8-33% faster than sk=4 - split_k = 8 - else: - # split_k=4 gives 2x more blocks than split_k=2 - split_k = 4 - elif use_splitk_for_large_hq: - # For h_q > 64 with bs > 8: - if total_topk >= 512: - # High topk: split_k=4 for all medium/large bs - split_k = 4 - else: - # Low topk: split_k=2 is sufficient - split_k = 2 - elif use_splitk_for_h64_large_topk: - split_k = 2 - else: - split_k = _select_split_k(total_topk, h_q, total_tokens) - topk_per_split = (total_topk + split_k - 1) // split_k # Get pre-allocated intermediate buffers @@ -2907,8 +2861,7 @@ def fused_gather_attn_decode_dsv4_dual_scope_low_overhead( stride_po = buffers["stride_po"] stride_plse = buffers["stride_plse"] - # Reuse pre-allocated output buffers to avoid torch.empty() calls - # that would be captured in CUDA graphs (each adds ~7-8us replay overhead). + # Allocate output buffers output = torch.empty(total_tokens, h_q, d_v, dtype=torch.bfloat16, device=device) lse = torch.empty(total_tokens, h_q, dtype=torch.float32, device=device) diff --git a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_optimized.py b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_optimized.py index 7426dd839..50abdafea 100644 --- a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_optimized.py +++ b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_optimized.py @@ -1,13 +1,13 @@ """ Optimized Triton MLA Decode Kernels for DeepSeek V4. -This module provides optimized sparse attention decode with reduced Python overhead. +This module provides optimized sparse attention decode. Key optimizations: 1. Fused gather+dequant+attention kernels (eliminates intermediate buffers) 2. Split-K for better GPU parallelism on small batches -3. Pre-allocated buffer pool for splitk intermediate results -4. Pre-computed strides to reduce tensor metadata operations +3. Proper dispatch: no-splitk for large batches, split-K for small batches +4. All paths use fused kernels (no 2-phase fallback) Note: This implementation assumes KV cache is always FP8 quantized. """ @@ -15,24 +15,38 @@ Note: This implementation assumes KV cache is always FP8 quantized. from typing import Optional, Tuple import torch -import triton -from .triton_mla_kernels_decode_common import ( - _bucket_total_tokens, - _unified_sparse_decode_kernel, - compute_token_ranges, -) -from .triton_mla_kernels_decode_dsv4 import ( - DSV4_D_QK, - fused_gather_dequant_fp8_dsv4, -) from .triton_mla_kernels_decode_fused import ( + DSV4_D_QK, fused_gather_attn_decode_dsv4, fused_gather_attn_decode_dsv4_dual_scope, fused_gather_attn_decode_dsv4_dual_scope_low_overhead, ) +def _should_use_fused_splitk(total_tokens: int, h_q: int, total_topk: int) -> bool: + """Determine whether to use fused split-K kernel (low overhead). + + The fused split-K kernel is preferred for small batch sizes because + split-K provides better GPU utilization when the grid is small. + + This matches the original _should_use_fused_dual_scope() thresholds. + """ + if total_tokens <= 4: + return True + if h_q <= 64 and total_topk <= 800: + return total_tokens <= 256 + if h_q <= 64 and total_topk >= 1024: + return total_tokens <= 128 + # h_q > 64 (e.g. h_q=128 when q is padded to full n_heads). + if h_q > 64: + if total_topk >= 400: + return total_tokens <= 32 + else: + return total_tokens <= 128 + return True + + def triton_sparse_attn_decode( q: torch.Tensor, kv_scope, @@ -54,67 +68,6 @@ def triton_sparse_attn_decode( ) -def _should_use_fused_dual_scope(total_tokens: int, h_q: int, total_topk: int) -> bool: - """Determine whether to use fused kernel for dual-scope cases. - - Returns True if the fused kernel (with splitk for small bs) should be used. - For large batch sizes (>= 256), use _should_use_fused_nosplitk instead. - - The thresholds below were determined empirically on MI355X (256 CUs). - """ - if total_tokens <= 4: - return True - if h_q <= 64 and total_topk <= 800: - return total_tokens <= 256 - if h_q <= 64 and total_topk >= 1024: - return total_tokens <= 128 - # h_q > 64 (e.g. h_q=128 when q is padded to full n_heads). - if h_q > 64: - if total_topk >= 400: - return total_tokens <= 32 - else: - return total_tokens <= 128 - return True - - -def _should_use_fused_nosplitk(total_tokens: int, h_q: int, total_topk: int) -> bool: - """Determine whether to use the fused no-splitk kernel for large batches. - - Kernel-level benchmarking on MI355X shows that for large batch sizes - (total_tokens >= 256), the fused dual-scope kernel WITHOUT split-K - is ~10% faster than the separate gather+attention path: - - total_tokens=256: fused-noSK=169us vs separate=194us (14% faster) - total_tokens=512: fused-noSK=350us vs separate=408us (14% faster) - total_tokens=1024: fused-noSK=700us vs separate=777us (10% faster) - total_tokens=4096: fused-noSK=2761us vs separate=3063us (10% faster) - - The fused no-splitk kernel avoids: - 1. Materializing the large intermediate gathered_kv buffer - 2. The separate gather kernel launch - 3. The split-K combine overhead - - For total_tokens < 256, the separate path is faster because the - fused kernel has insufficient parallelism. - - For extend (total_tokens >= 1024), the fused kernel always wins - regardless of h_q or total_topk because: - - The grid already has thousands of blocks (good GPU utilization) - - It eliminates 1.5-5 GB gathered_kv buffer allocation - - It eliminates 2x gather_dequant kernel launches (~414 us) - - It avoids chunking that TP>1 configs require with the separate path - """ - if total_tokens >= 1024: - return True - if h_q <= 64: - return False # Not benchmarked for h_q <= 64 - if total_topk < 200: - return False # Small topk doesn't benefit - # For h_q > 64 and total_topk >= 200: - # Fused no-splitk wins for total_tokens >= 256 - return total_tokens >= 256 - - def _triton_sparse_attn_decode_dsv4( q: torch.Tensor, kv_scope, @@ -123,10 +76,17 @@ def _triton_sparse_attn_decode_dsv4( d_v: int, attn_sink: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: - """Optimized sparse attention decode for DeepSeek V4 (d_qk=512).""" + """Sparse attention decode for DeepSeek V4 (d_qk=512). + + All paths use fused kernels (no 2-phase fallback). + + Dispatch logic: + - Single scope: always use fused kernel + - Dual scope, small total_tokens: fused split-K kernel (low overhead) + - Dual scope, otherwise: fused no-splitk kernel + """ b, s_q, h_q, d_qk = q.shape total_tokens = b * s_q - device = q.device topk_main = kv_scope.indices_in_kvcache.shape[-1] kv_quantized_main = kv_scope.blocked_k_quantized @@ -134,52 +94,55 @@ def _triton_sparse_attn_decode_dsv4( # Single scope case if extra_kv_scope is None: - if topk_main < 8192: - q_reshaped = q.reshape(total_tokens, h_q, d_qk) - if not q_reshaped.is_contiguous(): - q_reshaped = q_reshaped.contiguous() - - indices_main = kv_scope.indices_in_kvcache.reshape(total_tokens, topk_main) - if not indices_main.is_contiguous(): - indices_main = indices_main.contiguous() - - output, lse = fused_gather_attn_decode_dsv4( - q_reshaped, - kv_quantized_main, - indices_main, - block_size_main, - sm_scale, - topk_length=kv_scope.topk_length, - attn_sink=attn_sink, - s_q=s_q, - ) - return output.view(b, s_q, h_q, d_v), lse.view(b, s_q, h_q).transpose(1, 2) - else: - from .triton_mla_kernels_decode_dsv4 import triton_sparse_attn_decode_dsv4 - - return triton_sparse_attn_decode_dsv4( - q, kv_scope, extra_kv_scope, sm_scale, d_v, attn_sink - ) - - # Dual scope case - topk_extra = extra_kv_scope.indices_in_kvcache.shape[-1] - total_topk = topk_main + topk_extra - - # For large batch sizes, use fused no-splitk kernel (10% faster than separate). - # This check is BEFORE the chunking check because the fused kernel does NOT - # allocate the intermediate gathered_kv buffer, so buffer size limits don't apply. - if _should_use_fused_nosplitk(total_tokens, h_q, total_topk): q_reshaped = q.reshape(total_tokens, h_q, d_qk).contiguous() - indices_main = kv_scope.indices_in_kvcache.reshape( total_tokens, topk_main ).contiguous() - block_size_extra = extra_kv_scope.blocked_k.shape[1] - indices_extra = extra_kv_scope.indices_in_kvcache.reshape( - total_tokens, topk_extra - ).contiguous() + output, lse = fused_gather_attn_decode_dsv4( + q_reshaped, + kv_quantized_main, + indices_main, + block_size_main, + sm_scale, + topk_length=kv_scope.topk_length, + attn_sink=attn_sink, + s_q=s_q, + ) + return output.view(b, s_q, h_q, d_v), lse.view(b, s_q, h_q).transpose(1, 2) + # Dual scope case + topk_extra = extra_kv_scope.indices_in_kvcache.shape[-1] + total_topk = topk_main + topk_extra + block_size_extra = extra_kv_scope.blocked_k.shape[1] + + q_reshaped = q.reshape(total_tokens, h_q, d_qk).contiguous() + indices_main = kv_scope.indices_in_kvcache.reshape( + total_tokens, topk_main + ).contiguous() + indices_extra = extra_kv_scope.indices_in_kvcache.reshape( + total_tokens, topk_extra + ).contiguous() + + # Dispatch: use split-K for small batches, no-splitk for everything else. + if _should_use_fused_splitk(total_tokens, h_q, total_topk): + # Small batch: fused split-K kernel (better GPU utilization) + output, lse = fused_gather_attn_decode_dsv4_dual_scope_low_overhead( + q_reshaped, + kv_quantized_main, + indices_main, + block_size_main, + extra_kv_scope.blocked_k_quantized, + indices_extra, + block_size_extra, + sm_scale, + topk_length_main=kv_scope.topk_length, + topk_length_extra=extra_kv_scope.topk_length, + attn_sink=attn_sink, + s_q=s_q, + ) + else: + # Large batch / extend / prefill: fused no-splitk kernel output, lse = fused_gather_attn_decode_dsv4_dual_scope( q_reshaped, kv_quantized_main, @@ -195,153 +158,5 @@ def _triton_sparse_attn_decode_dsv4( s_q=s_q, force_no_splitk=True, ) - return output.view(b, s_q, h_q, d_v), lse.view(b, s_q, h_q).transpose(1, 2) - - # Check if chunking needed for separate path (fall back to original implementation) - token_ranges = compute_token_ranges(total_tokens, total_topk, d_qk) - if len(token_ranges) > 1: - from .triton_mla_kernels_decode_dsv4 import triton_sparse_attn_decode_dsv4 - - return triton_sparse_attn_decode_dsv4( - q, kv_scope, extra_kv_scope, sm_scale, d_v, attn_sink - ) - - # Use fused dual-scope kernel with low-overhead buffer pool - if _should_use_fused_dual_scope(total_tokens, h_q, total_topk): - q_reshaped = q.reshape(total_tokens, h_q, d_qk).contiguous() - - indices_main = kv_scope.indices_in_kvcache.reshape( - total_tokens, topk_main - ).contiguous() - - block_size_extra = extra_kv_scope.blocked_k.shape[1] - indices_extra = extra_kv_scope.indices_in_kvcache.reshape( - total_tokens, topk_extra - ).contiguous() - - output, lse = fused_gather_attn_decode_dsv4_dual_scope_low_overhead( - q_reshaped, - kv_quantized_main, - indices_main, - block_size_main, - extra_kv_scope.blocked_k_quantized, - indices_extra, - block_size_extra, - sm_scale, - topk_length_main=kv_scope.topk_length, - topk_length_extra=extra_kv_scope.topk_length, - attn_sink=attn_sink, - s_q=s_q, - ) - return output.view(b, s_q, h_q, d_v), lse.view(b, s_q, h_q).transpose(1, 2) - - # Fallback: Separate gather + attention path - return _fallback_gather_attention( - q, - kv_scope, - extra_kv_scope, - sm_scale, - d_v, - attn_sink, - total_tokens, - h_q, - d_qk, - topk_main, - topk_extra, - block_size_main, - kv_quantized_main, - fused_gather_dequant_fp8_dsv4, - ) - - -def _fallback_gather_attention( - q: torch.Tensor, - kv_scope, - extra_kv_scope, - sm_scale: float, - d_v: int, - attn_sink: Optional[torch.Tensor], - total_tokens: int, - h_q: int, - d_qk: int, - topk_main: int, - topk_extra: int, - block_size_main: int, - kv_quantized_main, - fused_gather_fn, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Fallback path: separate gather + attention kernels.""" - b = q.shape[0] - s_q = q.shape[1] - device = q.device - total_topk = topk_main + topk_extra - - gathered_kv = torch.empty( - total_tokens, total_topk, d_qk, dtype=torch.bfloat16, device=device - ) - invalid_mask = torch.empty( - total_tokens, total_topk, dtype=torch.bool, device=device - ) - output = torch.empty(total_tokens, h_q, d_v, dtype=torch.bfloat16, device=device) - lse = torch.empty(total_tokens, h_q, dtype=torch.float32, device=device) - - indices_main = kv_scope.indices_in_kvcache.reshape(total_tokens, topk_main) - block_size_extra = extra_kv_scope.blocked_k.shape[1] - indices_extra = extra_kv_scope.indices_in_kvcache.reshape(total_tokens, topk_extra) - - fused_gather_fn( - kv_quantized_main, - indices_main, - block_size_main, - kv_scope.topk_length, - extra_kv_scope.blocked_k_quantized, - indices_extra, - block_size_extra, - extra_kv_scope.topk_length, - gathered_kv, - invalid_mask, - s_q, - ) - - if q.dtype == torch.bfloat16 and q.is_contiguous(): - q_reshaped = q.view(total_tokens, h_q, d_qk) - else: - q_reshaped = q.to(torch.bfloat16).reshape(total_tokens, h_q, d_qk) - if not q_reshaped.is_contiguous(): - q_reshaped = q_reshaped.contiguous() - - HAS_ATTN_SINK = attn_sink is not None - attn_sink_tensor = attn_sink if HAS_ATTN_SINK else lse[:1] - - grid = lambda meta: (total_tokens, triton.cdiv(h_q, meta["BLOCK_H"])) - _unified_sparse_decode_kernel[grid]( - q_reshaped, - gathered_kv, - invalid_mask, - attn_sink_tensor, - output, - lse, - sm_scale, - total_tokens, - _bucket_total_tokens(total_tokens), - h_q, - total_topk, - d_qk, - d_v, - q_reshaped.stride(0), - q_reshaped.stride(1), - q_reshaped.stride(2), - gathered_kv.stride(0), - gathered_kv.stride(1), - gathered_kv.stride(2), - invalid_mask.stride(0), - invalid_mask.stride(1), - output.stride(0), - output.stride(1), - output.stride(2), - lse.stride(0), - lse.stride(1), - HAS_ATTN_SINK=HAS_ATTN_SINK, - ) return output.view(b, s_q, h_q, d_v), lse.view(b, s_q, h_q).transpose(1, 2) diff --git a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_splitk.py b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_splitk.py deleted file mode 100644 index 2f6c6e789..000000000 --- a/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_splitk.py +++ /dev/null @@ -1,534 +0,0 @@ -""" -Split-K Attention Kernel for Large TopK Cases - -This module implements a split-K version of the attention kernel that: -1. Splits the K (topk) dimension across multiple kernel instances -2. Each instance computes partial results with its own m_i, l_i, and accumulators -3. A combine kernel merges the partial results using online softmax - -This reduces register pressure by processing fewer K tokens per kernel instance, -improving occupancy and overall performance for large topk cases. -""" - -from typing import Optional, Tuple - -import torch -import triton -import triton.language as tl - -from .triton_mla_kernels_decode_common import _bucket_total_tokens - - -# ============================================================================ -# Split-K Attention Kernel -# ============================================================================ -@triton.autotune( - configs=[ - # Split-K attention on already-gathered BF16 KV. - # - BLOCK_N=256: amortizes memory access over KV tokens (memory-bound kernel). - # - BLOCK_D=128: matches KV tile structure. - # - num_warps=8, num_stages=2: memory-bound kernel benefits from more warps - # and software pipelining (overlaps memory loads with compute). - # - BLOCK_H varies for different batch sizes: - triton.Config( - {"BLOCK_H": 16, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=2 - ), - triton.Config( - {"BLOCK_H": 32, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=2 - ), - triton.Config( - {"BLOCK_H": 64, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=2 - ), - triton.Config( - {"BLOCK_H": 128, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=2 - ), - ], - key=["total_tokens_bucket", "h_q", "topk_per_split", "d_qk"], -) -@triton.jit -def _splitk_attention_kernel( - Q, - KV, - Mask, - PartialOutput, - PartialLSE, - PartialM, - sm_scale, - total_tokens, - total_tokens_bucket, - h_q, - total_topk, - d_qk, - d_v, - topk_per_split, - stride_q_t, - stride_q_h, - stride_q_d, - stride_kv_t, - stride_kv_k, - stride_kv_d, - stride_mask_t, - stride_mask_k, - stride_po_s, - stride_po_t, - stride_po_h, - stride_po_d, - stride_plse_s, - stride_plse_t, - stride_plse_h, - stride_pm_s, - stride_pm_t, - stride_pm_h, - BLOCK_H: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_D: tl.constexpr, -): - """Split-K attention kernel that processes a subset of K tokens.""" - LOG2E: tl.constexpr = 1.4426950408889634 - - pid_t = tl.program_id(0) - pid_h = tl.program_id(1) - pid_k = tl.program_id(2) - pid_t_64 = pid_t.to(tl.int64) - - NEG_INF = float("-inf") - - offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) - mask_h = offs_h < h_q - - # Compute K range for this split - k_start = pid_k * topk_per_split - k_end = tl.minimum(k_start + topk_per_split, total_topk) - - m_i = tl.full([BLOCK_H], NEG_INF, dtype=tl.float32) - l_i = tl.zeros([BLOCK_H], dtype=tl.float32) - - acc_0 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - acc_1 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - acc_2 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - acc_3 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - - stride_q_t_64 = tl.cast(stride_q_t, tl.int64) - stride_kv_t_64 = tl.cast(stride_kv_t, tl.int64) - stride_mask_t_64 = tl.cast(stride_mask_t, tl.int64) - q_base = Q + pid_t_64 * stride_q_t_64 - kv_base = KV + pid_t_64 * stride_kv_t_64 - mask_base = Mask + pid_t_64 * stride_mask_t_64 - - for n_start in range(k_start, k_end, BLOCK_N): - offs_n = n_start + tl.arange(0, BLOCK_N) - mask_n = offs_n < k_end - - mask_ptrs = mask_base + offs_n * stride_mask_k - invalid = tl.load(mask_ptrs, mask=mask_n, other=True) - valid = mask_n & ~invalid - - qk = tl.zeros([BLOCK_H, BLOCK_N], dtype=tl.float32) - - for d_start in range(0, d_qk, BLOCK_D): - offs_d = d_start + tl.arange(0, BLOCK_D) - mask_d = offs_d < d_qk - - q_ptrs = ( - q_base + offs_h[:, None] * stride_q_h + offs_d[None, :] * stride_q_d - ) - q_chunk = tl.load( - q_ptrs, mask=mask_h[:, None] & mask_d[None, :], other=0.0 - ).to(tl.bfloat16) - - k_ptrs = ( - kv_base + offs_n[:, None] * stride_kv_k + offs_d[None, :] * stride_kv_d - ) - k_chunk = tl.load( - k_ptrs, mask=valid[:, None] & mask_d[None, :], other=0.0 - ).to(tl.bfloat16) - - qk += tl.dot(q_chunk, tl.trans(k_chunk)) - - qk = qk * sm_scale - qk = tl.where(valid[None, :], qk, NEG_INF) - - m_ij = tl.max(qk, axis=1) - m_new = tl.maximum(m_i, m_ij) - alpha = tl.where(m_i == NEG_INF, 0.0, tl.math.exp2((m_i - m_new) * LOG2E)) - p = tl.where(qk == NEG_INF, 0.0, tl.math.exp2((qk - m_new[:, None]) * LOG2E)) - l_new = alpha * l_i + tl.sum(p, axis=1) - p_bf16 = p.to(tl.bfloat16) - - offs_v = tl.arange(0, BLOCK_D) - v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d - v = tl.load(v_ptrs, mask=valid[:, None], other=0.0).to(tl.bfloat16) - acc_0 = acc_0 * alpha[:, None] + tl.dot(p_bf16, v) - - offs_v = BLOCK_D + tl.arange(0, BLOCK_D) - v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d - v = tl.load( - v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0 - ).to(tl.bfloat16) - acc_1 = acc_1 * alpha[:, None] + tl.dot(p_bf16, v) - - offs_v = 2 * BLOCK_D + tl.arange(0, BLOCK_D) - v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d - v = tl.load( - v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0 - ).to(tl.bfloat16) - acc_2 = acc_2 * alpha[:, None] + tl.dot(p_bf16, v) - - offs_v = 3 * BLOCK_D + tl.arange(0, BLOCK_D) - v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d - v = tl.load( - v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0 - ).to(tl.bfloat16) - acc_3 = acc_3 * alpha[:, None] + tl.dot(p_bf16, v) - - m_i = m_new - l_i = l_new - - # Store partial results - stride_po_s_64 = tl.cast(stride_po_s, tl.int64) - stride_po_t_64 = tl.cast(stride_po_t, tl.int64) - po_base = PartialOutput + pid_k * stride_po_s_64 + pid_t_64 * stride_po_t_64 - - offs_h_2d = offs_h[:, None] - mask_h_2d = mask_h[:, None] - offs_v_0 = tl.arange(0, BLOCK_D) - offs_v_1 = BLOCK_D + tl.arange(0, BLOCK_D) - offs_v_2 = 2 * BLOCK_D + tl.arange(0, BLOCK_D) - offs_v_3 = 3 * BLOCK_D + tl.arange(0, BLOCK_D) - - tl.store( - po_base + offs_h_2d * stride_po_h + offs_v_0[None, :] * stride_po_d, - acc_0, - mask=mask_h_2d, - ) - tl.store( - po_base + offs_h_2d * stride_po_h + offs_v_1[None, :] * stride_po_d, - acc_1, - mask=mask_h_2d & (offs_v_1[None, :] < d_v), - ) - tl.store( - po_base + offs_h_2d * stride_po_h + offs_v_2[None, :] * stride_po_d, - acc_2, - mask=mask_h_2d & (offs_v_2[None, :] < d_v), - ) - tl.store( - po_base + offs_h_2d * stride_po_h + offs_v_3[None, :] * stride_po_d, - acc_3, - mask=mask_h_2d & (offs_v_3[None, :] < d_v), - ) - - stride_plse_s_64 = tl.cast(stride_plse_s, tl.int64) - stride_plse_t_64 = tl.cast(stride_plse_t, tl.int64) - plse_ptrs = ( - PartialLSE - + pid_k * stride_plse_s_64 - + pid_t_64 * stride_plse_t_64 - + offs_h * stride_plse_h - ) - tl.store(plse_ptrs, l_i, mask=mask_h) - - stride_pm_s_64 = tl.cast(stride_pm_s, tl.int64) - stride_pm_t_64 = tl.cast(stride_pm_t, tl.int64) - pm_ptrs = ( - PartialM - + pid_k * stride_pm_s_64 - + pid_t_64 * stride_pm_t_64 - + offs_h * stride_pm_h - ) - tl.store(pm_ptrs, m_i, mask=mask_h) - - -# ============================================================================ -# Combine Kernel for Split-K -# ============================================================================ -@triton.autotune( - configs=[ - # Simple reduce kernel merging split-K results. - # - BLOCK_D=128: 4 iterations to cover d_v=512. - # - num_warps=4: sufficient for this simple reduce operation. - # - BLOCK_H varies for different batch sizes: - triton.Config({"BLOCK_H": 16, "BLOCK_D": 128}, num_warps=4, num_stages=1), - triton.Config({"BLOCK_H": 32, "BLOCK_D": 128}, num_warps=4, num_stages=1), - triton.Config({"BLOCK_H": 64, "BLOCK_D": 128}, num_warps=4, num_stages=1), - ], - key=["total_tokens_bucket", "h_q", "split_k"], -) -@triton.jit -def _combine_splitk_attention_kernel( - PartialOutput, - PartialLSE, - PartialM, - AttnSink, - Output, - LSE, - total_tokens, - total_tokens_bucket, - h_q, - d_v, - split_k, - stride_po_s, - stride_po_t, - stride_po_h, - stride_po_d, - stride_plse_s, - stride_plse_t, - stride_plse_h, - stride_pm_s, - stride_pm_t, - stride_pm_h, - stride_o_t, - stride_o_h, - stride_o_d, - stride_lse_t, - stride_lse_h, - HAS_ATTN_SINK: tl.constexpr, - BLOCK_H: tl.constexpr, - BLOCK_D: tl.constexpr, -): - """Combine partial results from split-K attention kernel.""" - LOG2E: tl.constexpr = 1.4426950408889634 - NEG_INF = float("-inf") - POS_INF = float("+inf") - - pid_t = tl.program_id(0) - pid_h = tl.program_id(1) - pid_t_64 = pid_t.to(tl.int64) - - offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) - mask_h = offs_h < h_q - - m_acc = tl.full([BLOCK_H], NEG_INF, dtype=tl.float32) - l_acc = tl.zeros([BLOCK_H], dtype=tl.float32) - - acc_0 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - acc_1 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - acc_2 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - acc_3 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32) - - stride_po_s_64 = tl.cast(stride_po_s, tl.int64) - stride_po_t_64 = tl.cast(stride_po_t, tl.int64) - stride_plse_s_64 = tl.cast(stride_plse_s, tl.int64) - stride_plse_t_64 = tl.cast(stride_plse_t, tl.int64) - stride_pm_s_64 = tl.cast(stride_pm_s, tl.int64) - stride_pm_t_64 = tl.cast(stride_pm_t, tl.int64) - - offs_h_2d = offs_h[:, None] - mask_h_2d = mask_h[:, None] - offs_v_0 = tl.arange(0, BLOCK_D) - offs_v_1 = BLOCK_D + tl.arange(0, BLOCK_D) - offs_v_2 = 2 * BLOCK_D + tl.arange(0, BLOCK_D) - offs_v_3 = 3 * BLOCK_D + tl.arange(0, BLOCK_D) - - for k in range(split_k): - k_64 = tl.cast(k, tl.int64) - po_base = PartialOutput + k_64 * stride_po_s_64 + pid_t_64 * stride_po_t_64 - - p_acc_0 = tl.load( - po_base + offs_h_2d * stride_po_h + offs_v_0[None, :] * stride_po_d, - mask=mask_h_2d, - other=0.0, - ) - p_acc_1 = tl.load( - po_base + offs_h_2d * stride_po_h + offs_v_1[None, :] * stride_po_d, - mask=mask_h_2d & (offs_v_1[None, :] < d_v), - other=0.0, - ) - p_acc_2 = tl.load( - po_base + offs_h_2d * stride_po_h + offs_v_2[None, :] * stride_po_d, - mask=mask_h_2d & (offs_v_2[None, :] < d_v), - other=0.0, - ) - p_acc_3 = tl.load( - po_base + offs_h_2d * stride_po_h + offs_v_3[None, :] * stride_po_d, - mask=mask_h_2d & (offs_v_3[None, :] < d_v), - other=0.0, - ) - - plse_ptrs = ( - PartialLSE - + k_64 * stride_plse_s_64 - + pid_t_64 * stride_plse_t_64 - + offs_h * stride_plse_h - ) - p_l = tl.load(plse_ptrs, mask=mask_h, other=0.0) - - pm_ptrs = ( - PartialM - + k_64 * stride_pm_s_64 - + pid_t_64 * stride_pm_t_64 - + offs_h * stride_pm_h - ) - p_m = tl.load(pm_ptrs, mask=mask_h, other=NEG_INF) - - m_new = tl.maximum(m_acc, p_m) - alpha_acc = tl.where( - m_acc == NEG_INF, 0.0, tl.math.exp2((m_acc - m_new) * LOG2E) - ) - alpha_p = tl.where(p_m == NEG_INF, 0.0, tl.math.exp2((p_m - m_new) * LOG2E)) - l_new = alpha_acc * l_acc + alpha_p * p_l - - acc_0 = acc_0 * alpha_acc[:, None] + p_acc_0 * alpha_p[:, None] - acc_1 = acc_1 * alpha_acc[:, None] + p_acc_1 * alpha_p[:, None] - acc_2 = acc_2 * alpha_acc[:, None] + p_acc_2 * alpha_p[:, None] - acc_3 = acc_3 * alpha_acc[:, None] + p_acc_3 * alpha_p[:, None] - - m_acc = m_new - l_acc = l_new - - lse = m_acc + tl.math.log2(tl.where(l_acc == 0.0, 1.0, l_acc)) / LOG2E - is_lonely_q = l_acc == 0.0 - - if HAS_ATTN_SINK: - attn_sink_vals = tl.load(AttnSink + offs_h, mask=mask_h, other=0.0) - exp_attn_sink_minus_m = tl.math.exp2((attn_sink_vals - m_acc) * LOG2E) - denominator = l_acc + exp_attn_sink_minus_m - denominator = tl.where(denominator == 0.0, 1.0, denominator) - output_scale = 1.0 / denominator - else: - output_scale = tl.where(l_acc == 0.0, 0.0, 1.0 / l_acc) - - is_lonely_q_2d = is_lonely_q[:, None] - output_scale_2d = output_scale[:, None] - acc_0 = tl.where(is_lonely_q_2d, 0.0, acc_0 * output_scale_2d) - acc_1 = tl.where(is_lonely_q_2d, 0.0, acc_1 * output_scale_2d) - acc_2 = tl.where(is_lonely_q_2d, 0.0, acc_2 * output_scale_2d) - acc_3 = tl.where(is_lonely_q_2d, 0.0, acc_3 * output_scale_2d) - lse = tl.where(is_lonely_q, POS_INF, lse) - - stride_o_t_64 = tl.cast(stride_o_t, tl.int64) - o_base = Output + pid_t_64 * stride_o_t_64 - - tl.store( - o_base + offs_h_2d * stride_o_h + offs_v_0[None, :] * stride_o_d, - acc_0.to(tl.bfloat16), - mask=mask_h_2d, - ) - tl.store( - o_base + offs_h_2d * stride_o_h + offs_v_1[None, :] * stride_o_d, - acc_1.to(tl.bfloat16), - mask=mask_h_2d & (offs_v_1[None, :] < d_v), - ) - tl.store( - o_base + offs_h_2d * stride_o_h + offs_v_2[None, :] * stride_o_d, - acc_2.to(tl.bfloat16), - mask=mask_h_2d & (offs_v_2[None, :] < d_v), - ) - tl.store( - o_base + offs_h_2d * stride_o_h + offs_v_3[None, :] * stride_o_d, - acc_3.to(tl.bfloat16), - mask=mask_h_2d & (offs_v_3[None, :] < d_v), - ) - - stride_lse_t_64 = tl.cast(stride_lse_t, tl.int64) - tl.store(LSE + pid_t_64 * stride_lse_t_64 + offs_h * stride_lse_h, lse, mask=mask_h) - - -# ============================================================================ -# Runner Function -# ============================================================================ -def run_splitk_attention( - q_reshaped: torch.Tensor, - gathered_kv: torch.Tensor, - invalid_mask: torch.Tensor, - d_v: int, - sm_scale: float, - total_tokens: int, - h_q: int, - total_topk: int, - d_qk: int, - attn_sink: Optional[torch.Tensor] = None, - split_k: int = 4, -) -> Tuple[torch.Tensor, torch.Tensor]: - """Run split-K attention kernel.""" - device = q_reshaped.device - - topk_per_split = (total_topk + split_k - 1) // split_k - - partial_output = torch.empty( - split_k, total_tokens, h_q, d_v, dtype=torch.float32, device=device - ) - partial_lse = torch.empty( - split_k, total_tokens, h_q, dtype=torch.float32, device=device - ) - partial_m = torch.empty( - split_k, total_tokens, h_q, dtype=torch.float32, device=device - ) - - output = torch.empty(total_tokens, h_q, d_v, dtype=torch.bfloat16, device=device) - lse = torch.empty(total_tokens, h_q, dtype=torch.float32, device=device) - - grid_splitk = lambda meta: ( - total_tokens, - triton.cdiv(h_q, meta["BLOCK_H"]), - split_k, - ) - _splitk_attention_kernel[grid_splitk]( - q_reshaped, - gathered_kv, - invalid_mask, - partial_output, - partial_lse, - partial_m, - sm_scale, - total_tokens, - _bucket_total_tokens(total_tokens), - h_q, - total_topk, - d_qk, - d_v, - topk_per_split, - q_reshaped.stride(0), - q_reshaped.stride(1), - q_reshaped.stride(2), - gathered_kv.stride(0), - gathered_kv.stride(1), - gathered_kv.stride(2), - invalid_mask.stride(0), - invalid_mask.stride(1), - partial_output.stride(0), - partial_output.stride(1), - partial_output.stride(2), - partial_output.stride(3), - partial_lse.stride(0), - partial_lse.stride(1), - partial_lse.stride(2), - partial_m.stride(0), - partial_m.stride(1), - partial_m.stride(2), - ) - - HAS_ATTN_SINK = attn_sink is not None - attn_sink_tensor = attn_sink if HAS_ATTN_SINK else lse[:1] - - grid_combine = lambda meta: (total_tokens, triton.cdiv(h_q, meta["BLOCK_H"])) - _combine_splitk_attention_kernel[grid_combine]( - partial_output, - partial_lse, - partial_m, - attn_sink_tensor, - output, - lse, - total_tokens, - _bucket_total_tokens(total_tokens), - h_q, - d_v, - split_k, - partial_output.stride(0), - partial_output.stride(1), - partial_output.stride(2), - partial_output.stride(3), - partial_lse.stride(0), - partial_lse.stride(1), - partial_lse.stride(2), - partial_m.stride(0), - partial_m.stride(1), - partial_m.stride(2), - output.stride(0), - output.stride(1), - output.stride(2), - lse.stride(0), - lse.stride(1), - HAS_ATTN_SINK=HAS_ATTN_SINK, - ) - - return output, lse