From 3f5e2c76882510a44f1ec31914371a5d9b406762 Mon Sep 17 00:00:00 2001 From: kk <43161300+kkHuang-amd@users.noreply.github.com> Date: Tue, 26 May 2026 14:54:40 +0800 Subject: [PATCH] [AMD] Dsv4/pr2 compressor opt (#26208) Co-authored-by: wunhuang Co-authored-by: Thomas Wang <1am9trash@gmail.com> Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com> Co-authored-by: HaiShaw Co-authored-by: amd-danli103 Co-authored-by: Lin, Soga Co-authored-by: Raiden-Makoto Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com> Co-authored-by: yichiche@amd.com Co-authored-by: yctseng0211 Co-authored-by: Bingxu Chen --- docs/diffusion/compatibility_matrix.md | 2 - python/sglang/srt/environ.py | 4 + python/sglang/srt/layers/activation.py | 14 + .../deepseek_v4_backend_hip_radix.py | 18 +- .../srt/layers/attention/dsv4/compress_hip.py | 9 +- .../srt/layers/attention/dsv4/compressor.py | 38 +- .../layers/attention/dsv4/compressor_v2.py | 543 ++- .../attention/dsv4/fused_compress_triton.py | 954 +++++ .../srt/layers/attention/dsv4/indexer.py | 202 +- .../srt/layers/attention/dsv4/metadata.py | 8 +- .../srt/layers/attention/hip_flash_mla.py | 15 +- .../attention/nsa/triton_decode/__init__.py | 98 + .../triton_mla_kernels_decode_common.py | 585 ++++ .../triton_mla_kernels_decode_dsv4.py | 1355 ++++++++ .../triton_mla_kernels_decode_fused.py | 3089 +++++++++++++++++ .../triton_mla_kernels_decode_optimized.py | 289 ++ .../triton_mla_kernels_decode_splitk.py | 534 +++ python/sglang/srt/layers/deepseek_v4_rope.py | 86 + python/sglang/srt/layers/fused_qk_norm.py | 157 + .../sglang/srt/layers/moe/moe_runner/aiter.py | 5 + python/sglang/srt/layers/moe/topk.py | 52 +- python/sglang/srt/layers/quantization/fp8.py | 23 +- .../srt/mem_cache/deepseek_v4_memory_pool.py | 7 +- python/sglang/srt/models/deepseek_v2.py | 7 +- python/sglang/srt/models/deepseek_v4.py | 162 +- sgl-kernel/benchmark/bench_dsv4_norm_rope.py | 75 + sgl-kernel/csrc/common_extension.cc | 15 + sgl-kernel/include/sgl_kernel_ops.h | 29 + sgl-kernel/python/sgl_kernel/__init__.py | 8 + sgl-kernel/tests/test_dsv4_norm_rope.py | 130 + .../dsv4/test_fused_compress_attn_hip.py | 465 +++ 31 files changed, 8829 insertions(+), 149 deletions(-) create mode 100644 python/sglang/srt/layers/attention/dsv4/fused_compress_triton.py create mode 100644 python/sglang/srt/layers/attention/nsa/triton_decode/__init__.py create mode 100644 python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_common.py create mode 100644 python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_dsv4.py create mode 100644 python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_fused.py create mode 100644 python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_optimized.py create mode 100644 python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_splitk.py create mode 100644 python/sglang/srt/layers/fused_qk_norm.py create mode 100644 sgl-kernel/benchmark/bench_dsv4_norm_rope.py create mode 100644 sgl-kernel/tests/test_dsv4_norm_rope.py create mode 100644 test/manual/dsv4/test_fused_compress_attn_hip.py diff --git a/docs/diffusion/compatibility_matrix.md b/docs/diffusion/compatibility_matrix.md index 435a83c36..cd7f2d5b5 100644 --- a/docs/diffusion/compatibility_matrix.md +++ b/docs/diffusion/compatibility_matrix.md @@ -66,8 +66,6 @@ default parameters when initializing and generating videos. | FLUX.2-dev-NVFP4 | `black-forest-labs/FLUX.2-dev-NVFP4` | | FLUX.2-Klein-4B | `black-forest-labs/FLUX.2-klein-4B` | | FLUX.2-Klein-9B | `black-forest-labs/FLUX.2-klein-9B` | -| FLUX.2-Klein-Base-4B | `black-forest-labs/FLUX.2-klein-base-4B` | -| FLUX.2-Klein-Base-9B | `black-forest-labs/FLUX.2-klein-base-9B` | | Z-Image | `Tongyi-MAI/Z-Image` | | Z-Image-Turbo | `Tongyi-MAI/Z-Image-Turbo` | | GLM-Image | `zai-org/GLM-Image` | diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 85eb45e09..28a5813ae 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -627,7 +627,9 @@ class Envs: SGLANG_OPT_USE_TRITON_SWA_PREPARE = EnvBool(True) SGLANG_OPT_USE_AITER_MHC_PRE = EnvBool(True) SGLANG_OPT_USE_AITER_MHC_POST = EnvBool(True) + SGLANG_OPT_USE_AITER_SILU_MUL = EnvBool(False) SGLANG_OPT_USE_FUSED_COMPRESS = EnvBool(False) + SGLANG_OPT_USE_FUSED_COMPRESS_TRITON = EnvBool(False) SGLANG_OPT_USE_FUSED_QK_NORM_ROPE = EnvBool(True) SGLANG_OPT_USE_FUSED_CLAMP_ACT_MUL = EnvBool(True) SGLANG_FIX_MTP_HC_HIDDEN = EnvBool(False) @@ -644,6 +646,7 @@ class Envs: SGLANG_OPT_USE_TILELANG_MHC_PRE = EnvBool(True) SGLANG_OPT_USE_TILELANG_MHC_POST = EnvBool(True) SGLANG_OPT_USE_TILELANG_INDEXER = EnvBool(False) + SGLANG_OPT_USE_AITER_INDEXER = EnvBool(False) SGLANG_OPT_USE_JIT_INDEXER_METADATA = EnvBool(True) SGLANG_OPT_USE_ONLINE_COMPRESS = EnvBool(False) SGLANG_OPT_USE_COMPRESSOR_V2 = EnvBool(True) @@ -688,6 +691,7 @@ class Envs: # Cache / overlap SGLANG_OPT_USE_FUSED_STORE_CACHE = EnvBool(True) + SGLANG_OPT_USE_JIT_NORM = EnvBool(True) SGLANG_OPT_USE_MULTI_STREAM_OVERLAP = EnvBool(True) # CUDA graph diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index 216e37a23..a76c454fa 100644 --- a/python/sglang/srt/layers/activation.py +++ b/python/sglang/srt/layers/activation.py @@ -33,6 +33,7 @@ from sglang.srt.layers.utils import MultiPlatformOp from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( cpu_has_amx_support, + get_bool_env_var, is_cpu, is_cuda, is_hip, @@ -50,6 +51,7 @@ _is_cpu_amx_available = cpu_has_amx_support() _is_cpu = is_cpu() _is_hip = is_hip() _is_xpu = is_xpu() +_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _is_cuda: from sglang.jit_kernel.activation import ( @@ -71,6 +73,9 @@ elif _is_musa: return torch.empty(output_shape, dtype=x.dtype, device=x.device) +if _use_aiter: + from aiter import silu_and_mul as _aiter_silu_and_mul + if is_npu(): import torch_npu @@ -82,6 +87,8 @@ class SiluAndMul(MultiPlatformOp): super().__init__(*args, **kwargs) if get_global_server_args().rl_on_policy_target is not None: self._forward_method = self.forward_native + elif _use_aiter and envs.SGLANG_OPT_USE_AITER_SILU_MUL.get(): + self._forward_method = self.forward_aiter def forward_native(self, x: torch.Tensor) -> torch.Tensor: d = x.shape[-1] // 2 @@ -94,6 +101,13 @@ class SiluAndMul(MultiPlatformOp): silu_and_mul(x, out) return out + def forward_aiter(self, x: torch.Tensor, limit: float = 0.0) -> torch.Tensor: + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + _aiter_silu_and_mul(out, x, limit) + return out + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: if _is_cpu_amx_available: out = torch.ops.sgl_kernel.silu_and_mul_cpu(x) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 033b877ba..3e0ee41ab 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -20,11 +20,19 @@ import torch.nn.functional as F from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.dsv4.compressor import ( - CompressorBackendMixin, - FusedCompressMetadata, - create_paged_compressor_data, -) + +if envs.SGLANG_OPT_USE_COMPRESSOR_V2.get(): + from sglang.srt.layers.attention.dsv4.compressor_v2 import ( + CompressorBackendMixin, + FusedCompressMetadata, + create_paged_compressor_data, + ) +else: + from sglang.srt.layers.attention.dsv4.compressor import ( + CompressorBackendMixin, + FusedCompressMetadata, + create_paged_compressor_data, + ) from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin from sglang.srt.layers.attention.dsv4.metadata import ( PagedIndexerMetadata, diff --git a/python/sglang/srt/layers/attention/dsv4/compress_hip.py b/python/sglang/srt/layers/attention/dsv4/compress_hip.py index cff010c55..8c6b7df9b 100644 --- a/python/sglang/srt/layers/attention/dsv4/compress_hip.py +++ b/python/sglang/srt/layers/attention/dsv4/compress_hip.py @@ -12,9 +12,14 @@ import triton.language as tl from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa.dsa_indexer import rotate_activation from sglang.srt.layers.attention.dsv4.compressor import Compressor as _CompressorBase +from sglang.srt.layers.attention.dsv4.fused_compress_triton import ( + fused_ape_pool_norm_rope, +) +from sglang.srt.layers.attention.nsa.nsa_indexer import rotate_activation from sglang.srt.layers.deepseek_v4_rope import ( apply_rotary_emb_triton, fused_norm_rope_inplace_triton, + fused_softmax_pool_triton, ) try: @@ -372,10 +377,6 @@ class CompressorHip(_CompressorBase): freqs_real_table = self._get_freqs_cis_real() freqs_batch = freqs_real_table[comp_positions] - from sglang.srt.layers.attention.dsv4.fused_compress_kernel import ( - fused_ape_pool_norm_rope, - ) - kv_compressed = fused_ape_pool_norm_rope( kv_score_gathered=gathered, ape=self.ape, diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index e663c6ab8..fa326592f 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -57,6 +57,9 @@ class CompressorBackendMixin: assert isinstance(metadata, FusedCompressMetadata) return metadata + def _maybe_upgrade_forward_metadata(self) -> None: + pass + def forward_compress( self, *, @@ -91,6 +94,37 @@ class CompressorBackendMixin: metadata = (forward_batch.req_pool_indices.to(torch.int32), None, plan) indices, extra_data, plan = metadata + if _is_hip: + if not is_paged: + raise NotImplementedError("HIP fused compressor expects paged metadata") + + from sglang.srt.layers.attention.dsv4.fused_compress_triton import ( + hip_compress_forward, + hip_compress_fused_norm_rope_inplace, + ) + + kv_compressed = hip_compress_forward( + kv_score_buffer=kv_score_buffer, + kv_score_input=kv_score_input, + ape=ape, + indices=indices, + plan=plan, + compress_ratio=compress_ratio, + head_dim=head_dim, + extra_data=extra_data, + ) + norm_eps = ( + norm.variance_epsilon if hasattr(norm, "variance_epsilon") else norm.eps + ) + hip_compress_fused_norm_rope_inplace( + kv_compressed, + norm.weight, + norm_eps, + freqs_cis_cache, + plan, + ) + return rotate_activation(kv_compressed) if rotate else kv_compressed + kv_compressed = compress_forward( kv_score_buffer=kv_score_buffer, kv_score_input=kv_score_input, @@ -279,6 +313,8 @@ def create_paged_compressor_data( if is_overlap: write_overlap_loc = get_raw_loc(write_positions - compress_ratio) extra_data = write_overlap_loc.view(-1, 1) + elif _is_hip: + extra_data = get_raw_loc(write_positions - compress_ratio) else: extra_data = None plan = CompressorDecodePlan(compress_ratio, seq_lens.to(torch.int32)) @@ -392,7 +428,7 @@ class Compressor(nn.Module): ) -if _is_hip: +if _is_hip and not envs.SGLANG_OPT_USE_COMPRESSOR_V2.get(): from sglang.srt.layers.attention.dsv4.compress_hip import ( # noqa: F811 CompressorHip as Compressor, ) diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py index 13b5c65cc..41063e3a8 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py @@ -10,6 +10,7 @@ from sglang.jit_kernel.dsv4 import ( compress_forward, compress_norm_rope_store, ) +from sglang.jit_kernel.utils import is_hip_runtime from sglang.srt.environ import envs if TYPE_CHECKING: @@ -24,12 +25,380 @@ CompressMetadata: TypeAlias = Union[CompressorDecodePlan, CompressorPrefillPlan] # NOTE: alias for backward compatibility FusedCompressMetadata: TypeAlias = CompressMetadata +_is_hip = is_hip_runtime() + +if _is_hip: + import triton + import triton.language as tl + + @triton.jit + def _c128_compress_decode_kernel( + buf_ptr, + input_ptr, + ape_ptr, + out_ptr, + plan_ptr, + buf_stride_slot, + input_stride_b, + ape_stride_r, + out_stride_b, + bs, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + ): + """Fused C128 decode: write to state buffer + online softmax-pool. + + plan_ptr points to int32 view: [bs, 4] where each row is + {seq_len, write_loc, read_page_0, read_page_1}. + """ + bid = tl.program_id(0) + if bid >= bs: + return + + # Parse plan + plan_base = plan_ptr + bid * 4 + seq_len = tl.load(plan_base).to(tl.int32) + write_loc = tl.load(plan_base + 1).to(tl.int32) + read_page_0 = tl.load(plan_base + 2).to(tl.int32) + + d = tl.arange(0, BLOCK_D) + last_dim: tl.constexpr = HEAD_DIM * 2 + + # Step 1: Write kv_score_input to state buffer at write_loc + d_mask_full = d < last_dim + input_val = tl.load( + input_ptr + bid * input_stride_b + d, mask=d_mask_full, other=0.0 + ) + tl.store(buf_ptr + write_loc * buf_stride_slot + d, input_val, mask=d_mask_full) + + # Step 2: Check boundary condition + d_mask_hd = d < HEAD_DIM + if seq_len % COMPRESS_RATIO != 0: + tl.store( + out_ptr + bid * out_stride_b + d, + tl.zeros([BLOCK_D], tl.float32), + mask=d_mask_hd, + ) + return + + # Step 3: Online softmax-pool over 128 slots in the page + page_base = read_page_0 * COMPRESS_RATIO * buf_stride_slot + m_prev = tl.full([BLOCK_D], float("-inf"), tl.float32) + kv_acc = tl.zeros([BLOCK_D], tl.float32) + w_acc = tl.zeros([BLOCK_D], tl.float32) + + for k in tl.static_range(COMPRESS_RATIO): + slot_addr = page_base + k * buf_stride_slot + kv_val = tl.load(buf_ptr + slot_addr + d, mask=d_mask_hd, other=0.0).to( + tl.float32 + ) + sc_val = tl.load( + buf_ptr + slot_addr + HEAD_DIM + d, mask=d_mask_hd, other=0.0 + ).to(tl.float32) + ape_val = tl.load( + ape_ptr + k * ape_stride_r + d, mask=d_mask_hd, other=0.0 + ).to(tl.float32) + score_k = sc_val + ape_val + + m_new = tl.maximum(m_prev, score_k) + exp_old = tl.where(m_prev == float("-inf"), 0.0, tl.exp(m_prev - m_new)) + exp_cur = tl.where(score_k == float("-inf"), 0.0, tl.exp(score_k - m_new)) + kv_acc = kv_acc * exp_old + exp_cur * kv_val + w_acc = w_acc * exp_old + exp_cur + m_prev = m_new + + compressed = kv_acc / w_acc + tl.store(out_ptr + bid * out_stride_b + d, compressed, mask=d_mask_hd) + + @triton.jit + def _c128_compress_prefill_write_kernel( + buf_ptr, + input_ptr, + plan_w_ptr, + buf_stride_slot, + input_stride_b, + num_w, + BLOCK_D: tl.constexpr, + LAST_DIM: tl.constexpr, + ): + """Prefill write phase: scatter kv_score_input tokens into state buffer.""" + wid = tl.program_id(0) + if wid >= num_w: + return + + # WritePlan: {ragged_id(u32), write_loc(i32)} = 8 bytes = 2 int32s + plan_base = plan_w_ptr + wid * 2 + ragged_id = (tl.load(plan_base).to(tl.int32)) & 0xFFFF + write_loc = tl.load(plan_base + 1).to(tl.int32) + + d = tl.arange(0, BLOCK_D) + d_mask = d < LAST_DIM + + if write_loc >= 0: + input_val = tl.load( + input_ptr + ragged_id * input_stride_b + d, mask=d_mask, other=0.0 + ) + tl.store(buf_ptr + write_loc * buf_stride_slot + d, input_val, mask=d_mask) + + @triton.jit + def _c128_compress_prefill_compress_kernel( + buf_ptr, + ape_ptr, + out_ptr, + plan_c_ptr, + buf_stride_slot, + ape_stride_r, + out_stride_b, + num_c, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + ): + """Prefill compress phase: online softmax-pool for each compress plan entry.""" + cid = tl.program_id(0) + if cid >= num_c: + return + + # CompressPlan: {seq_len(u32), ragged_id(u16)|buffer_len(u16), read_page_0(i32), read_page_1(i32)} + plan_base = plan_c_ptr + cid * 4 + read_page_0 = tl.load(plan_base + 2).to(tl.int32) + + d = tl.arange(0, BLOCK_D) + d_mask_hd = d < HEAD_DIM + + if read_page_0 < 0: + tl.store( + out_ptr + cid * out_stride_b + d, + tl.zeros([BLOCK_D], tl.float32), + mask=d_mask_hd, + ) + return + + page_base = read_page_0 * COMPRESS_RATIO * buf_stride_slot + m_prev = tl.full([BLOCK_D], float("-inf"), tl.float32) + kv_acc = tl.zeros([BLOCK_D], tl.float32) + w_acc = tl.zeros([BLOCK_D], tl.float32) + + for k in tl.static_range(COMPRESS_RATIO): + slot_addr = page_base + k * buf_stride_slot + kv_val = tl.load(buf_ptr + slot_addr + d, mask=d_mask_hd, other=0.0).to( + tl.float32 + ) + sc_val = tl.load( + buf_ptr + slot_addr + HEAD_DIM + d, mask=d_mask_hd, other=0.0 + ).to(tl.float32) + ape_val = tl.load( + ape_ptr + k * ape_stride_r + d, mask=d_mask_hd, other=0.0 + ).to(tl.float32) + score_k = sc_val + ape_val + + m_new = tl.maximum(m_prev, score_k) + exp_old = tl.where(m_prev == float("-inf"), 0.0, tl.exp(m_prev - m_new)) + exp_cur = tl.where(score_k == float("-inf"), 0.0, tl.exp(score_k - m_new)) + kv_acc = kv_acc * exp_old + exp_cur * kv_val + w_acc = w_acc * exp_old + exp_cur + m_prev = m_new + + compressed = kv_acc / w_acc + tl.store(out_ptr + cid * out_stride_b + d, compressed, mask=d_mask_hd) + + +def _compress_forward_c128_triton( + kv_score_buffer: torch.Tensor, + kv_score_input: torch.Tensor, + ape: torch.Tensor, + plan: Union[CompressorDecodePlan, CompressorPrefillPlan], + head_dim: int, +) -> torch.Tensor: + """Triton C128 compress_forward for HIP (wave64). + + Fuses write + online-softmax-pool into Triton kernels. + CUDA graph compatible. + """ + num_total_slots = kv_score_buffer.shape[0] * kv_score_buffer.shape[1] + num_pages = kv_score_buffer.shape[0] + last_dim = kv_score_buffer.shape[-1] + compress_ratio = 128 + + buf_flat = kv_score_buffer.view(-1, last_dim) + buf_stride_slot = last_dim # elements per slot + + BLOCK_D = triton.next_power_of_2(last_dim) + + if plan.is_decode: + # Decode path: single kernel does write + compress + plan_raw = plan[1].view(torch.int32) # [bs, 4] + bs = plan_raw.shape[0] + out = torch.empty( + bs, head_dim, dtype=torch.float32, device=kv_score_input.device + ) + + if bs > 0 and num_total_slots > 0: + grid = (bs,) + _c128_compress_decode_kernel[grid]( + buf_flat, + kv_score_input, + ape, + out, + plan_raw, + buf_stride_slot, + kv_score_input.stride(0), + ape.stride(0), + out.stride(0), + bs, + HEAD_DIM=head_dim, + BLOCK_D=triton.next_power_of_2(head_dim), + COMPRESS_RATIO=compress_ratio, + num_warps=8, + ) + return out + else: + # Prefill path: separate write kernel + compress kernel + plan_c_raw = plan[1].view(torch.int32) # [num_c, 4] + plan_w = plan[2] # [num_w, 8] uint8 + plan_w_raw = plan_w.view(torch.int32) # [num_w, 2] + num_c = plan_c_raw.shape[0] + num_w = plan_w_raw.shape[0] + + out = torch.empty( + num_c, head_dim, dtype=torch.float32, device=kv_score_input.device + ) + + # Phase 1: Write + if num_w > 0 and num_total_slots > 0: + grid_w = (num_w,) + _c128_compress_prefill_write_kernel[grid_w]( + buf_flat, + kv_score_input, + plan_w_raw, + buf_stride_slot, + kv_score_input.stride(0), + num_w, + BLOCK_D=BLOCK_D, + LAST_DIM=last_dim, + num_warps=4, + ) + + # Phase 2: Compress + if num_c > 0 and num_pages > 0: + grid_c = (num_c,) + _c128_compress_prefill_compress_kernel[grid_c]( + buf_flat, + ape, + out, + plan_c_raw, + buf_stride_slot, + ape.stride(0), + out.stride(0), + num_c, + HEAD_DIM=head_dim, + BLOCK_D=triton.next_power_of_2(head_dim), + COMPRESS_RATIO=compress_ratio, + num_warps=8, + ) + + return out + def _use_online_compress(compress_ratio: int) -> bool: """Online state-pool path is c128-only.""" return compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get() +def _extract_positions_from_plan( + plan: Union[CompressorDecodePlan, CompressorPrefillPlan], + compress_ratio: int, +) -> torch.Tensor: + """Extract RoPE positions from plan tensors (decode or prefill). + + DecodePlan layout: [bs, 16] uint8, first 4 bytes = uint32 seq_len. + CompressPlan layout: [num_c, 16] uint8, first 4 bytes = uint32 seq_len. + Position for RoPE = seq_len - compress_ratio. + """ + plan_tensor = plan[1] # plan_d or plan_c + seq_lens = plan_tensor[:, :4].contiguous().view(torch.int32).squeeze(-1) + positions = seq_lens.to(torch.int32) - compress_ratio + return positions + + +def _compress_forward_c128_fallback( + kv_score_buffer: torch.Tensor, + kv_score_input: torch.Tensor, + ape: torch.Tensor, + plan: Union[CompressorDecodePlan, CompressorPrefillPlan], + head_dim: int, +) -> torch.Tensor: + """PyTorch fallback for C128 compress_forward on HIP (wave64). + + Fully vectorized, compatible with CUDA graph capture. + kv_score_buffer: [num_pages, 128, head_dim * 2] + ape: [128, head_dim] + + IMPORTANT: This also performs the write to state buffer (like the JIT kernel). + The JIT kernel does: (1) write kv_score_input to buffer, (2) compress from buffer. + """ + num_total_slots = kv_score_buffer.shape[0] * kv_score_buffer.shape[1] + num_pages = kv_score_buffer.shape[0] + last_dim = kv_score_buffer.shape[-1] + + # Step 1: WRITE kv_score_input to state buffer + if num_total_slots > 0: + buf_flat = kv_score_buffer.view(-1, last_dim) + if plan.is_decode: + # Decode: plan_d has write_loc per batch item + plan_raw = plan[1].view(torch.int32) # [bs, 4] + write_locs = plan_raw[:, 1].long() + # Only write valid locations (>= 0 and < buffer size) + valid_write = (write_locs >= 0) & (write_locs < num_total_slots) + if valid_write.any(): + buf_flat[write_locs[valid_write]] = kv_score_input[valid_write] + else: + # Prefill: plan_w has {ragged_id, write_loc} per write entry + plan_w = plan[2] # [num_w, 8] uint8 = WritePlan + if plan_w.shape[0] > 0: + plan_w_raw = plan_w.view(torch.int32) # [num_w, 2] + ragged_ids = plan_w_raw[:, 0].long() & 0xFFFF + write_locs = plan_w_raw[:, 1].long() + valid_write = (write_locs >= 0) & (write_locs < num_total_slots) + ragged_ids_safe = ragged_ids.clamp( + min=0, max=kv_score_input.shape[0] - 1 + ) + if valid_write.any(): + buf_flat[write_locs[valid_write]] = kv_score_input[ + ragged_ids_safe[valid_write] + ] + + # Step 2: COMPRESS (read from buffer page and do softmax-pool) + plan_c = plan[1] # plan_d for decode, plan_c for prefill + num_tokens = plan_c.shape[0] + if num_pages == 0 or num_tokens == 0: + return kv_score_input.new_zeros(num_tokens, head_dim) + + plan_c_raw = plan_c.view(torch.int32) # [N, 4] + read_page_0 = plan_c_raw[:, 2].long() + # Use torch.where instead of clamp to handle -1 (invalid) gracefully + valid_read = (read_page_0 >= 0) & (read_page_0 < num_pages) + read_page_0_safe = torch.where( + valid_read, read_page_0, torch.zeros_like(read_page_0) + ) + + gathered = kv_score_buffer[read_page_0_safe] # [N, 128, head_dim*2] + kv = gathered[:, :, :head_dim].float() + score = gathered[:, :, head_dim:].float() + ape.float().unsqueeze(0) + weights = score.softmax(dim=1) + out = (weights * kv).sum(dim=1) + + # For decode: zero out non-boundary tokens (seq_len % 128 != 0) + # so they don't corrupt kvcache location 0 when stored. + if plan.is_decode: + seq_lens = plan_c_raw[:, 0].to(torch.int32) + is_boundary = (seq_lens % 128 == 0).unsqueeze(-1) # [N, 1] + out = torch.where(is_boundary, out, torch.zeros_like(out)) + + return out.to(kv_score_input.dtype) + + class CompressorBackendMixin: def __init__(self): super().__init__() @@ -74,6 +443,8 @@ class CompressorBackendMixin: last_dim = 2 * head_dim * coff assert kv_score_buffer.shape[-1] == last_dim kv_score_buffer = kv_score_buffer.view(-1, compress_ratio, last_dim) + + # Step 1: compress_forward kv_compressed = compress_forward( kv_score_buffer=kv_score_buffer, kv_score_input=kv_score_input, @@ -83,7 +454,8 @@ class CompressorBackendMixin: head_dim=head_dim, is_online=is_online, ) - # NOTE: we use some hack here... + + # Step 2: norm + rope + store compress_norm_rope_store( kv_compressed, plan, @@ -109,34 +481,153 @@ class CompressorBackendMixin: token_to_kv_pool = self.token_to_kv_pool token_to_kv_pool = cast("DeepSeekV4TokenToKVPool", token_to_kv_pool) kv_score_input = compressor.compute_kv_score(x, forward_batch) + state_pool = compressor.get_state_pool(self) - out_loc = self._get_out_loc(compressor.ratio) - if compressor.is_in_indexer: - kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id) - page_size = token_to_kv_pool.get_index_k_page_size() + if _is_hip and not envs.SGLANG_OPT_USE_JIT_NORM.get(): + self._forward_unified_hip( + token_to_kv_pool=token_to_kv_pool, + kv_score_input=kv_score_input, + state_pool=state_pool, + compressor=compressor, + layer_id=layer_id, + ) else: - _, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id] - assert compress_kv_pool is not None - kv_cache = token_to_kv_pool.get_extra_key_buffer(layer_id) - page_size = token_to_kv_pool.get_extra_key_page_size(layer_id) - if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"): - # The v2 compressor writes directly into the raw C4 KV tensor. - # HiSparse C4 therefore needs the physical C4 location here. - out_loc = compress_kv_pool.translate_loc_to_hisparse_device(out_loc) - self._forward_compress_all_in_one( - kv_score_buffer=state_pool.kv_score_buffer.kv_score, - kv_score_input=kv_score_input, - ape=compressor.ape, - head_dim=compressor.head_dim, - norm=compressor.norm, - freqs_cis_cache=compressor.freqs_cis, - kv_cache=kv_cache.view(dtype=torch.uint8), - is_indexer=compressor.is_in_indexer, - rotate=compressor.rotate, - compress_ratio=compressor.ratio, - page_size=page_size, - out_loc=out_loc, + out_loc = self._get_out_loc(compressor.ratio) + if compressor.is_in_indexer: + kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id) + page_size = token_to_kv_pool.get_index_k_page_size() + else: + _, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id] + assert compress_kv_pool is not None + kv_cache = token_to_kv_pool.get_extra_key_buffer(layer_id) + page_size = token_to_kv_pool.get_extra_key_page_size(layer_id) + if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"): + # The v2 compressor writes directly into the raw C4 KV tensor. + # HiSparse C4 therefore needs the physical C4 location here. + out_loc = compress_kv_pool.translate_loc_to_hisparse_device(out_loc) + self._forward_compress_all_in_one( + kv_score_buffer=state_pool.kv_score_buffer.kv_score, + kv_score_input=kv_score_input, + ape=compressor.ape, + head_dim=compressor.head_dim, + norm=compressor.norm, + freqs_cis_cache=compressor.freqs_cis, + kv_cache=kv_cache.view(dtype=torch.uint8), + is_indexer=compressor.is_in_indexer, + rotate=compressor.rotate, + compress_ratio=compressor.ratio, + page_size=page_size, + out_loc=out_loc, + ) + + def _forward_unified_hip( + self, + token_to_kv_pool: DeepSeekV4TokenToKVPool, + kv_score_input: torch.Tensor, + state_pool, + compressor: Compressor, + layer_id: int, + ) -> None: + """HIP-specific forward path using PyTorch/Triton fallbacks.""" + from sglang.srt.layers.attention.dsv4.quant_k_cache import ( + quant_to_nope_fp8_rope_bf16_pack_triton, ) + from sglang.srt.layers.attention.nsa.nsa_indexer import rotate_activation + from sglang.srt.layers.attention.nsa.triton_kernel import act_quant + from sglang.srt.layers.deepseek_v4_rope import fused_norm_rope_inplace_triton + + compress_ratio = compressor.ratio + head_dim = compressor.head_dim + is_indexer = compressor.is_in_indexer + + plan = self._get_paged_compress_metadata(compress_ratio) + out_loc = self._get_out_loc(compress_ratio) + + # Step 1: compress_forward (always use JIT for both C4 and C128) + coff = 2 if is_overlap_compress(compress_ratio) else 1 + last_dim = 2 * head_dim * coff + kv_score_buffer = state_pool.kv_score_buffer.kv_score + kv_score_buffer = kv_score_buffer.view(-1, compress_ratio, last_dim) + + kv_compressed = compress_forward( + kv_score_buffer=kv_score_buffer, + kv_score_input=kv_score_input, + ape=compressor.ape.view(-1, head_dim), + plan=plan, + compress_ratio=compress_ratio, + head_dim=head_dim, + is_online=False, + ) + + if kv_compressed.shape[0] == 0: + return + + # For decode: zero out non-boundary tokens to prevent corrupting kvcache loc 0. + if plan.is_decode: + plan_raw = plan[1].view(torch.int32) + seq_lens_plan = plan_raw[:, 0].to(torch.int32) + is_boundary = (seq_lens_plan % compress_ratio == 0).unsqueeze(-1) + kv_compressed = torch.where( + is_boundary, kv_compressed, torch.zeros_like(kv_compressed) + ) + + # Step 2: norm + rope (Triton fallback for precision parity with V1) + positions = _extract_positions_from_plan(plan, compress_ratio) + positions_safe = positions.clamp(min=0) + + fused_norm_rope_inplace_triton( + kv_compressed, + compressor.norm.weight, + compressor.norm.variance_epsilon, + compressor.freqs_cis, + positions=positions_safe, + ) + + # Step 3: optional Hadamard rotation for indexer + if compressor.rotate: + kv_compressed = rotate_activation(kv_compressed) + + # Step 4: store to kvcache + # For decode: store ALL tokens. Non-boundary tokens have out_loc=0 (safe). + # For prefill: plan_c already only contains valid entries. + if plan.is_decode: + kv_to_store = kv_compressed + out_loc_to_store = out_loc + else: + kv_to_store = kv_compressed + plan_raw = plan[1].view(torch.int32) + ragged_ids = plan_raw[:, 1].to(torch.int32) & 0xFFFF + out_loc_to_store = out_loc[ragged_ids.long()] + + if kv_to_store.shape[0] == 0: + return + + if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get(): + # fused kernel: BF16 in -> FP8 quant + paged scatter in one launch + if is_indexer: + token_to_kv_pool.set_index_k_fused( + layer_id=layer_id, + loc=out_loc_to_store, + cache_k=kv_to_store, + ) + else: + token_to_kv_pool.set_extra_key_buffer_fused( + layer_id=layer_id, + loc=out_loc_to_store, + cache_k=kv_to_store, + ) + else: + if is_indexer: + kv_fp8, kv_scale = act_quant(kv_to_store) + token_to_kv_pool.set_index_k_scale_buffer( + layer_id=layer_id, + loc=out_loc_to_store, + index_k=kv_fp8, + index_k_scale=kv_scale, + ) + else: + pack = quant_to_nope_fp8_rope_bf16_pack_triton(kv_to_store.bfloat16()) + token_to_kv_pool.set_extra_key_buffer(layer_id, out_loc_to_store, pack) # NOTE: alias for backward compatibility forward_indexer_compressor = forward_unified diff --git a/python/sglang/srt/layers/attention/dsv4/fused_compress_triton.py b/python/sglang/srt/layers/attention/dsv4/fused_compress_triton.py new file mode 100644 index 000000000..9434556c4 --- /dev/null +++ b/python/sglang/srt/layers/attention/dsv4/fused_compress_triton.py @@ -0,0 +1,954 @@ +"""HIP fused compressor kernels using the NV/main metadata contract. + +The public wrappers mirror ``compress_forward``: + + decode: indices, seq_lens, extra_data + prefill: indices, compress_plan, write_plan, extra_data + +Prefill plans are the upstream 16-byte ``PrefillPlan`` structs stored as +``uint8[:, 16]``. The wrappers reinterpret them as ``int32[:, 4]`` before +launching Triton kernels. +""" + +from __future__ import annotations + +from typing import Optional, Union + +import torch +import triton +import triton.language as tl + +from sglang.jit_kernel.dsv4.compress_old import ( + CompressorDecodePlan, + CompressorPrefillPlan, +) + + +@triton.jit +def _fused_ape_pool_norm_rope_kernel( + kv_score_ptr, + kv_score_stride_b, + kv_score_stride_k, + ape_ptr, + ape_stride_r, + rms_weight_ptr, + rms_eps, + freqs_ptr, + freqs_stride_b, + out_ptr, + out_stride_b, + head_dim, + rope_head_dim, + half_dim, + RATIO: tl.constexpr, + K_POOL: tl.constexpr, + BLOCK_D: tl.constexpr, + HALF_ROPE: tl.constexpr, + OVERLAP: tl.constexpr, +): + bid = tl.program_id(0) + d = tl.arange(0, BLOCK_D) + d_mask = d < head_dim + + m_prev = tl.full([BLOCK_D], float("-inf"), tl.float32) + kv_acc = tl.zeros([BLOCK_D], tl.float32) + w_acc = tl.zeros([BLOCK_D], tl.float32) + + batch_base = bid * kv_score_stride_b + + for k in tl.range(0, K_POOL): + if OVERLAP: + is_b = k >= RATIO + col_off = tl.where(is_b, head_dim, 0) + else: + col_off = 0 + + row_off = batch_base + k * kv_score_stride_k + kv_val = tl.load( + kv_score_ptr + row_off + col_off + d, mask=d_mask, other=0.0 + ).to(tl.float32) + sc_val = tl.load( + kv_score_ptr + row_off + half_dim + col_off + d, mask=d_mask, other=0.0 + ).to(tl.float32) + + ape_val = tl.load( + ape_ptr + (k % RATIO) * ape_stride_r + col_off + d, mask=d_mask, other=0.0 + ).to(tl.float32) + score_k = sc_val + ape_val + + m_new = tl.maximum(m_prev, score_k) + exp_old = tl.where(m_prev == float("-inf"), 0.0, tl.exp(m_prev - m_new)) + exp_cur = tl.where(score_k == float("-inf"), 0.0, tl.exp(score_k - m_new)) + kv_acc = kv_acc * exp_old + exp_cur * kv_val + w_acc = w_acc * exp_old + exp_cur + m_prev = m_new + + compressed = kv_acc / w_acc + rms_w = tl.load(rms_weight_ptr + d, mask=d_mask, other=0.0) + c_sq = tl.where(d_mask, compressed * compressed, 0.0) + var = tl.sum(c_sq, axis=0) / head_dim + normed = compressed * tl.rsqrt(var + rms_eps) * rms_w + + out_base = out_ptr + bid * out_stride_b + tl.store(out_base + d, normed.to(out_ptr.dtype.element_ty), mask=d_mask) + + rope_start = head_dim - rope_head_dim + p = tl.arange(0, HALF_ROPE) + pmask = p < (rope_head_dim // 2) + xr = tl.load(out_base + rope_start + 2 * p, mask=pmask, other=0.0).to(tl.float32) + xi = tl.load(out_base + rope_start + 2 * p + 1, mask=pmask, other=0.0).to( + tl.float32 + ) + + freq_base = bid * freqs_stride_b + fr = tl.load(freqs_ptr + freq_base + 2 * p, mask=pmask, other=1.0).to(tl.float32) + fi = tl.load(freqs_ptr + freq_base + 2 * p + 1, mask=pmask, other=0.0).to( + tl.float32 + ) + + tl.store( + out_base + rope_start + 2 * p, + (xr * fr - xi * fi).to(out_ptr.dtype.element_ty), + mask=pmask, + ) + tl.store( + out_base + rope_start + 2 * p + 1, + (xr * fi + xi * fr).to(out_ptr.dtype.element_ty), + mask=pmask, + ) + + +def fused_ape_pool_norm_rope( + kv_score_gathered: torch.Tensor, + ape: torch.Tensor, + rms_weight: torch.Tensor, + rms_eps: float, + freqs_cis_real: torch.Tensor, + head_dim: int, + rope_head_dim: int, + ratio: int, + overlap: bool, +) -> torch.Tensor: + """Fused APE-add + overlap-transform + softmax-pool + RMSNorm + RoPE.""" + coff = 2 if overlap else 1 + bs = kv_score_gathered.shape[0] + k_in = kv_score_gathered.shape[1] + last_dim = kv_score_gathered.shape[2] + half_dim = last_dim // 2 + assert k_in == ratio * coff, f"k_in={k_in} != ratio*coff={ratio}*{coff}" + + out = torch.empty( + bs, head_dim, dtype=torch.float32, device=kv_score_gathered.device + ) + if bs == 0: + return out + + block_d = triton.next_power_of_2(head_dim) + half_rope = triton.next_power_of_2(rope_head_dim // 2) + num_warps = 4 if head_dim <= 256 else 8 + + _fused_ape_pool_norm_rope_kernel[(bs,)]( + kv_score_gathered, + kv_score_gathered.stride(0), + kv_score_gathered.stride(1), + ape, + ape.stride(0), + rms_weight, + rms_eps, + freqs_cis_real, + freqs_cis_real.stride(0), + out, + out.stride(0), + head_dim, + rope_head_dim, + half_dim, + RATIO=ratio, + K_POOL=k_in, + BLOCK_D=block_d, + HALF_ROPE=half_rope, + OVERLAP=int(overlap), + num_warps=num_warps, + ) + return out + + +@triton.jit +def _c4_decode_kernel( + kv_in_ptr, + out_ptr, + buffer_ptr, + ape_ptr, + indices_ptr, + seq_lens_ptr, + extra_ptr, + kv_in_row_stride, + out_row_stride, + buffer_page_stride, + buffer_slot_stride, + ape_row_stride, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, +): + bid = tl.program_id(0) + pid_d = tl.program_id(1) + d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + + index = tl.load(indices_ptr + bid).to(tl.int64) + index_prev = tl.load(extra_ptr + bid).to(tl.int64) + seq_len = tl.load(seq_lens_ptr + bid).to(tl.int32) + write_slot = (seq_len + 3) % 4 + + in_base = bid.to(tl.int64) * kv_in_row_stride + page_base = ( + index * buffer_page_stride + write_slot.to(tl.int64) * buffer_slot_stride + ) + + valid_index = index >= 0 + for ch in tl.static_range(4): + ch_off = ch * HEAD_DIM + val = tl.load(kv_in_ptr + in_base + ch_off + d_offs, mask=d_mask, other=0.0) + tl.store( + buffer_ptr + page_base + ch_off + d_offs, + val, + mask=d_mask & valid_index, + ) + + NEG_BIG: tl.constexpr = -1.0e9 + running_max = tl.full((BLOCK_D,), NEG_BIG, tl.float32) + running_sum = tl.zeros((BLOCK_D,), tl.float32) + weighted = tl.zeros((BLOCK_D,), tl.float32) + + for slot in tl.static_range(8): + if slot < 4: + page = index_prev + kv_off = 0 + score_off = 2 * HEAD_DIM + else: + page = index + kv_off = HEAD_DIM + score_off = 3 * HEAD_DIM + + src_pos = seq_len - 8 + slot + is_input = slot == 7 + write_pos = ((seq_len - 1) // 4) * 4 + page = tl.where(src_pos < write_pos, index_prev, index) + slot_in_page = src_pos % 4 + slot_base = ( + page * buffer_page_stride + slot_in_page.to(tl.int64) * buffer_slot_stride + ) + valid = src_pos >= 0 + if slot == 7: + kv = tl.load( + kv_in_ptr + in_base + kv_off + d_offs, + mask=d_mask & valid, + other=0.0, + ) + score = tl.load( + kv_in_ptr + in_base + score_off + d_offs, + mask=d_mask & valid, + other=NEG_BIG, + ) + else: + kv = tl.load( + buffer_ptr + slot_base + kv_off + d_offs, + mask=d_mask & valid, + other=0.0, + ) + score = tl.load( + buffer_ptr + slot_base + score_off + d_offs, + mask=d_mask & valid, + other=NEG_BIG, + ) + bias = tl.load(ape_ptr + slot * ape_row_stride + d_offs, mask=d_mask, other=0.0) + s = score + bias + new_max = tl.maximum(running_max, s) + factor = tl.exp(running_max - new_max) + e = tl.where(valid, tl.exp(s - new_max), 0.0) + running_sum = running_sum * factor + e + weighted = weighted * factor + kv * e + running_max = new_max + + tl.store( + out_ptr + bid.to(tl.int64) * out_row_stride + d_offs, + weighted / running_sum, + mask=d_mask, + ) + + +@triton.jit +def _c4_prefill_compress_kernel( + kv_in_ptr, + out_ptr, + buffer_ptr, + ape_ptr, + indices_ptr, + extra_ptr, + plan_ptr, + kv_in_row_stride, + out_row_stride, + buffer_page_stride, + buffer_slot_stride, + ape_row_stride, + plan_row_stride, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_p = tl.program_id(0) + pid_d = tl.program_id(1) + d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + + plan_base = plan_ptr + pid_p * plan_row_stride + ragged_id = tl.load(plan_base + 0).to(tl.int32) + batch_id = tl.load(plan_base + 1).to(tl.int32) + position = tl.load(plan_base + 2).to(tl.int32) + window_len = tl.load(plan_base + 3).to(tl.int32) + if ragged_id < 0: + return + + extra_base = extra_ptr + batch_id.to(tl.int64) * 4 + load_first_page = tl.load(extra_base + 0).to(tl.int64) + load_second_page = tl.load(extra_base + 1).to(tl.int64) + + NEG_BIG: tl.constexpr = -1.0e9 + running_max = tl.full((BLOCK_D,), NEG_BIG, tl.float32) + running_sum = tl.zeros((BLOCK_D,), tl.float32) + weighted = tl.zeros((BLOCK_D,), tl.float32) + + for slot in tl.static_range(8): + in_state = slot < window_len + if slot < 4: + page = tl.where(window_len <= 4, load_second_page, load_first_page) + kv_off = 0 + score_off = 2 * HEAD_DIM + slot_in_page = slot + else: + page = load_second_page + kv_off = HEAD_DIM + score_off = 3 * HEAD_DIM + slot_in_page = slot - 4 + + src_pos = position - 7 + slot + state_valid = in_state & (src_pos >= 0) + slot_base = page * buffer_page_stride + slot_in_page * buffer_slot_stride + in_row = ragged_id - (7 - slot) + in_row_safe = tl.where(in_state, 0, in_row) + in_base = in_row_safe.to(tl.int64) * kv_in_row_stride + + kv_state = tl.load( + buffer_ptr + slot_base + kv_off + d_offs, + mask=d_mask & state_valid, + other=0.0, + ) + score_state = tl.load( + buffer_ptr + slot_base + score_off + d_offs, + mask=d_mask & state_valid, + other=NEG_BIG, + ) + kv_input = tl.load( + kv_in_ptr + in_base + kv_off + d_offs, + mask=d_mask & (~in_state), + other=0.0, + ) + score_input = tl.load( + kv_in_ptr + in_base + score_off + d_offs, + mask=d_mask & (~in_state), + other=NEG_BIG, + ) + kv = tl.where(in_state, kv_state, kv_input) + score = tl.where(in_state, score_state, score_input) + bias = tl.load(ape_ptr + slot * ape_row_stride + d_offs, mask=d_mask, other=0.0) + + s = score + bias + new_max = tl.maximum(running_max, s) + factor = tl.exp(running_max - new_max) + e = tl.exp(s - new_max) + running_sum = running_sum * factor + e + weighted = weighted * factor + kv * e + running_max = new_max + + tl.store( + out_ptr + ragged_id.to(tl.int64) * out_row_stride + d_offs, + weighted / running_sum, + mask=d_mask, + ) + + +@triton.jit +def _c4_prefill_write_kernel( + kv_in_ptr, + buffer_ptr, + indices_ptr, + extra_ptr, + plan_ptr, + kv_in_row_stride, + buffer_page_stride, + buffer_slot_stride, + plan_row_stride, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_p = tl.program_id(0) + pid_d = tl.program_id(1) + d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + + plan_base = plan_ptr + pid_p * plan_row_stride + ragged_id = tl.load(plan_base + 0).to(tl.int32) + batch_id = tl.load(plan_base + 1).to(tl.int32) + position = tl.load(plan_base + 2).to(tl.int32) + if ragged_id < 0: + return + + extra_base = extra_ptr + batch_id.to(tl.int64) * 4 + write_first_page = tl.load(extra_base + 2).to(tl.int64) + last_position = tl.load(extra_base + 3).to(tl.int32) + write_second_page = tl.load(indices_ptr + batch_id).to(tl.int64) + page = tl.where(position < last_position, write_first_page, write_second_page) + slot = position % 4 + + in_base = ragged_id.to(tl.int64) * kv_in_row_stride + dst_base = page * buffer_page_stride + slot.to(tl.int64) * buffer_slot_stride + for ch in tl.static_range(4): + ch_off = ch * HEAD_DIM + val = tl.load(kv_in_ptr + in_base + ch_off + d_offs, mask=d_mask, other=0.0) + tl.store(buffer_ptr + dst_base + ch_off + d_offs, val, mask=d_mask) + + +@triton.jit +def _c128_decode_kernel( + kv_in_ptr, + out_ptr, + buffer_ptr, + ape_ptr, + indices_ptr, + seq_lens_ptr, + extra_ptr, + kv_in_row_stride, + out_row_stride, + buffer_page_stride, + buffer_slot_stride, + ape_row_stride, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, +): + bid = tl.program_id(0) + pid_d = tl.program_id(1) + d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + + index = tl.load(indices_ptr + bid).to(tl.int64) + index_prev = tl.load(extra_ptr + bid).to(tl.int64) + seq_len = tl.load(seq_lens_ptr + bid).to(tl.int32) + write_slot = (seq_len + 127) % 128 + in_base = bid.to(tl.int64) * kv_in_row_stride + dst_base = index * buffer_page_stride + write_slot.to(tl.int64) * buffer_slot_stride + + for ch in tl.static_range(2): + ch_off = ch * HEAD_DIM + val = tl.load(kv_in_ptr + in_base + ch_off + d_offs, mask=d_mask, other=0.0) + tl.store(buffer_ptr + dst_base + ch_off + d_offs, val, mask=d_mask) + + NEG_BIG: tl.constexpr = -1.0e9 + running_max = tl.full((BLOCK_D,), NEG_BIG, tl.float32) + running_sum = tl.zeros((BLOCK_D,), tl.float32) + weighted = tl.zeros((BLOCK_D,), tl.float32) + + for chunk_start in tl.static_range(0, 128, BLOCK_S): + slot_offs = chunk_start + tl.arange(0, BLOCK_S) + src_pos = seq_len - 128 + slot_offs + valid = src_pos >= 0 + is_input = slot_offs == 127 + write_pos = ((seq_len - 1) // 128) * 128 + pages = tl.where(src_pos < write_pos, index_prev, index) + slot_in_page = src_pos % 128 + slot_bases = ( + pages * buffer_page_stride + slot_in_page.to(tl.int64) * buffer_slot_stride + ) + kv_tile = tl.load( + buffer_ptr + slot_bases[:, None] + d_offs[None, :], + mask=valid[:, None] & (~is_input)[:, None] & d_mask[None, :], + other=0.0, + ) + score_tile = tl.load( + buffer_ptr + slot_bases[:, None] + HEAD_DIM + d_offs[None, :], + mask=valid[:, None] & (~is_input)[:, None] & d_mask[None, :], + other=NEG_BIG, + ) + kv_input_tile = tl.load( + kv_in_ptr + in_base + d_offs[None, :], + mask=valid[:, None] & is_input[:, None] & d_mask[None, :], + other=0.0, + ) + score_input_tile = tl.load( + kv_in_ptr + in_base + HEAD_DIM + d_offs[None, :], + mask=valid[:, None] & is_input[:, None] & d_mask[None, :], + other=NEG_BIG, + ) + kv_tile = tl.where(is_input[:, None], kv_input_tile, kv_tile) + score_tile = tl.where(is_input[:, None], score_input_tile, score_tile) + bias_tile = tl.load( + ape_ptr + slot_offs[:, None] * ape_row_stride + d_offs[None, :], + mask=d_mask[None, :], + other=0.0, + ) + s = score_tile + bias_tile + local_max = tl.max(s, axis=0) + new_max = tl.maximum(running_max, local_max) + exp_s = tl.exp(s - new_max[None, :]) + exp_s = tl.where(valid[:, None], exp_s, 0.0) + factor = tl.exp(running_max - new_max) + running_sum = running_sum * factor + tl.sum(exp_s, axis=0) + weighted = weighted * factor + tl.sum(kv_tile * exp_s, axis=0) + running_max = new_max + + tl.store( + out_ptr + bid.to(tl.int64) * out_row_stride + d_offs, + weighted / running_sum, + mask=d_mask, + ) + + +@triton.jit +def _c128_prefill_compress_kernel( + kv_in_ptr, + out_ptr, + buffer_ptr, + ape_ptr, + indices_ptr, + plan_ptr, + kv_in_row_stride, + out_row_stride, + buffer_page_stride, + buffer_slot_stride, + ape_row_stride, + plan_row_stride, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_S: tl.constexpr, +): + pid_p = tl.program_id(0) + pid_d = tl.program_id(1) + d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + + plan_base = plan_ptr + pid_p * plan_row_stride + ragged_id = tl.load(plan_base + 0).to(tl.int32) + batch_id = tl.load(plan_base + 1).to(tl.int32) + position = tl.load(plan_base + 2).to(tl.int32) + window_len = tl.load(plan_base + 3).to(tl.int32) + if ragged_id < 0: + return + + index = tl.load(indices_ptr + batch_id).to(tl.int64) + NEG_BIG: tl.constexpr = -1.0e9 + running_max = tl.full((BLOCK_D,), NEG_BIG, tl.float32) + running_sum = tl.zeros((BLOCK_D,), tl.float32) + weighted = tl.zeros((BLOCK_D,), tl.float32) + + for chunk_start in tl.static_range(0, 128, BLOCK_S): + slot_offs = chunk_start + tl.arange(0, BLOCK_S) + is_state = slot_offs < window_len + src_pos = position - 127 + slot_offs + state_valid = is_state & (src_pos >= 0) + slot_bases = ( + index * buffer_page_stride + slot_offs.to(tl.int64) * buffer_slot_stride + ) + in_rows = ragged_id - (127 - slot_offs) + in_rows_safe = tl.where(is_state, tl.zeros_like(in_rows), in_rows) + in_bases = in_rows_safe.to(tl.int64) * kv_in_row_stride + + kv_state = tl.load( + buffer_ptr + slot_bases[:, None] + d_offs[None, :], + mask=state_valid[:, None] & d_mask[None, :], + other=0.0, + ) + score_state = tl.load( + buffer_ptr + slot_bases[:, None] + HEAD_DIM + d_offs[None, :], + mask=state_valid[:, None] & d_mask[None, :], + other=NEG_BIG, + ) + kv_input = tl.load( + kv_in_ptr + in_bases[:, None] + d_offs[None, :], + mask=(~is_state)[:, None] & d_mask[None, :], + other=0.0, + ) + score_input = tl.load( + kv_in_ptr + in_bases[:, None] + HEAD_DIM + d_offs[None, :], + mask=(~is_state)[:, None] & d_mask[None, :], + other=NEG_BIG, + ) + kv_tile = tl.where(is_state[:, None], kv_state, kv_input) + score_tile = tl.where(is_state[:, None], score_state, score_input) + bias_tile = tl.load( + ape_ptr + slot_offs[:, None] * ape_row_stride + d_offs[None, :], + mask=d_mask[None, :], + other=0.0, + ) + + s = score_tile + bias_tile + local_max = tl.max(s, axis=0) + new_max = tl.maximum(running_max, local_max) + exp_s = tl.exp(s - new_max[None, :]) + # Keep input-path entries valid; only state-path entries need src_pos guard. + valid = state_valid | (~is_state) + exp_s = tl.where(valid[:, None], exp_s, 0.0) + factor = tl.exp(running_max - new_max) + running_sum = running_sum * factor + tl.sum(exp_s, axis=0) + weighted = weighted * factor + tl.sum(kv_tile * exp_s, axis=0) + running_max = new_max + + tl.store( + out_ptr + ragged_id.to(tl.int64) * out_row_stride + d_offs, + weighted / running_sum, + mask=d_mask, + ) + + +@triton.jit +def _c128_prefill_write_kernel( + kv_in_ptr, + buffer_ptr, + indices_ptr, + plan_ptr, + kv_in_row_stride, + buffer_page_stride, + buffer_slot_stride, + plan_row_stride, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, +): + pid_p = tl.program_id(0) + pid_d = tl.program_id(1) + d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) + d_mask = d_offs < HEAD_DIM + + plan_base = plan_ptr + pid_p * plan_row_stride + ragged_id = tl.load(plan_base + 0).to(tl.int32) + batch_id = tl.load(plan_base + 1).to(tl.int32) + position = tl.load(plan_base + 2).to(tl.int32) + if ragged_id < 0: + return + + index = tl.load(indices_ptr + batch_id).to(tl.int64) + slot = position % 128 + in_base = ragged_id.to(tl.int64) * kv_in_row_stride + dst_base = index * buffer_page_stride + slot.to(tl.int64) * buffer_slot_stride + + for ch in tl.static_range(2): + ch_off = ch * HEAD_DIM + val = tl.load(kv_in_ptr + in_base + ch_off + d_offs, mask=d_mask, other=0.0) + tl.store(buffer_ptr + dst_base + ch_off + d_offs, val, mask=d_mask) + + +@triton.jit +def _compress_norm_rope_kernel( + kv_ptr, + weight_ptr, + freqs_ptr, + handle_ptr, + eps, + kv_row_stride, + freqs_row_stride, + plan_row_stride, + HEAD_DIM: tl.constexpr, + ROPE_DIM: tl.constexpr, + HEAD_BLOCK: tl.constexpr, + ROPE_PAIR_BLOCK: tl.constexpr, + COMPRESS_RATIO: tl.constexpr, + IS_DECODE: tl.constexpr, +): + work_id = tl.program_id(0) + + if IS_DECODE: + row = work_id + seq_len = tl.load(handle_ptr + work_id).to(tl.int32) + position = ((seq_len - 1) // COMPRESS_RATIO) * COMPRESS_RATIO + else: + plan_base = handle_ptr + work_id * plan_row_stride + row = tl.load(plan_base + 0).to(tl.int32) + plan_position = tl.load(plan_base + 2).to(tl.int32) + if row < 0: + return + position = plan_position + 1 - COMPRESS_RATIO + + base = row.to(tl.int64) * kv_row_stride + offs = tl.arange(0, HEAD_BLOCK) + mask = offs < HEAD_DIM + x = tl.load(kv_ptr + base + offs, mask=mask, other=0.0).to(tl.float32) + w = tl.load(weight_ptr + offs, mask=mask, other=0.0).to(tl.float32) + rms_inv = tl.rsqrt(tl.sum(x * x, axis=0) / HEAD_DIM + eps) + x_normed = x * rms_inv * w + + rope_start: tl.constexpr = HEAD_DIM - ROPE_DIM + pair_offs = tl.arange(0, ROPE_PAIR_BLOCK) + pair_mask = pair_offs < (ROPE_DIM // 2) + x_real = tl.load( + kv_ptr + base + rope_start + 2 * pair_offs, + mask=pair_mask, + other=0.0, + ).to(tl.float32) + x_imag = tl.load( + kv_ptr + base + rope_start + 2 * pair_offs + 1, + mask=pair_mask, + other=0.0, + ).to(tl.float32) + w_real = tl.load( + weight_ptr + rope_start + 2 * pair_offs, + mask=pair_mask, + other=1.0, + ).to(tl.float32) + w_imag = tl.load( + weight_ptr + rope_start + 2 * pair_offs + 1, + mask=pair_mask, + other=1.0, + ).to(tl.float32) + x_real = x_real * rms_inv * w_real + x_imag = x_imag * rms_inv * w_imag + + freq_base = position.to(tl.int64) * freqs_row_stride + f_real = tl.load(freqs_ptr + freq_base + 2 * pair_offs, mask=pair_mask, other=0.0) + f_imag = tl.load( + freqs_ptr + freq_base + 2 * pair_offs + 1, + mask=pair_mask, + other=0.0, + ) + out_real = x_real * f_real - x_imag * f_imag + out_imag = x_real * f_imag + x_imag * f_real + + tl.store(kv_ptr + base + offs, x_normed, mask=mask & (offs < rope_start)) + tl.store(kv_ptr + base + rope_start + 2 * pair_offs, out_real, mask=pair_mask) + tl.store(kv_ptr + base + rope_start + 2 * pair_offs + 1, out_imag, mask=pair_mask) + + +def _plan_as_i32(plan: torch.Tensor) -> torch.Tensor: + assert plan.dtype == torch.uint8 and plan.dim() == 2 and plan.shape[1] == 16 + return plan.view(torch.int32).view(-1, 4) + + +def _block_d(head_dim: int) -> int: + return min(32, triton.next_power_of_2(head_dim)) + + +def _check_common( + kv_score_buffer: torch.Tensor, + kv_score_input: torch.Tensor, + out: torch.Tensor, + ape: torch.Tensor, + indices: torch.Tensor, + head_dim: int, + compress_ratio: int, +) -> None: + coff = 2 if compress_ratio == 4 else 1 + assert kv_score_input.is_cuda and kv_score_buffer.is_cuda + assert kv_score_input.dim() == 2 and kv_score_input.dtype == torch.float32 + assert kv_score_input.shape[1] == 2 * coff * head_dim + assert kv_score_buffer.dim() == 3 and kv_score_buffer.dtype == torch.float32 + assert kv_score_buffer.shape[1:] == (compress_ratio, 2 * coff * head_dim) + assert out.shape == (kv_score_input.shape[0], head_dim) + assert out.dtype == torch.float32 and out.is_cuda + assert ape.shape == (compress_ratio * coff, head_dim) + assert ape.dtype == torch.float32 and ape.is_cuda + assert indices.dtype == torch.int32 and indices.is_cuda + + +def _is_decode_plan(plan: Union[CompressorDecodePlan, CompressorPrefillPlan]) -> bool: + return isinstance(plan, CompressorDecodePlan) + + +def hip_compress_forward( + *, + kv_score_buffer: torch.Tensor, + kv_score_input: torch.Tensor, + ape: torch.Tensor, + indices: torch.Tensor, + plan: Union[CompressorDecodePlan, CompressorPrefillPlan], + extra_data: Optional[torch.Tensor], + head_dim: int, + compress_ratio: int, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + if compress_ratio not in (4, 128): + raise ValueError(f"unsupported {compress_ratio=}") + if out is None: + out = kv_score_input.new_empty((kv_score_input.shape[0], head_dim)) + is_decode = _is_decode_plan(plan) + if not is_decode: + out.fill_(10000.0) + + _check_common( + kv_score_buffer, + kv_score_input, + out, + ape, + indices, + head_dim, + compress_ratio, + ) + + BLOCK_D = _block_d(head_dim) + num_d_chunks = triton.cdiv(head_dim, BLOCK_D) + + if is_decode: + seq_lens = plan.seq_lens + assert seq_lens.dtype == torch.int32 and seq_lens.is_cuda + assert seq_lens.shape == indices.shape + grid = (seq_lens.numel(), num_d_chunks) + if compress_ratio == 4: + assert extra_data is not None + assert extra_data.shape == (seq_lens.numel(), 1) + _c4_decode_kernel[grid]( + kv_score_input, + out, + kv_score_buffer, + ape, + indices, + seq_lens, + extra_data, + kv_score_input.stride(0), + out.stride(0), + kv_score_buffer.stride(0), + kv_score_buffer.stride(1), + ape.stride(0), + HEAD_DIM=head_dim, + BLOCK_D=BLOCK_D, + ) + else: + assert extra_data is not None + assert extra_data.shape == seq_lens.shape + _c128_decode_kernel[grid]( + kv_score_input, + out, + kv_score_buffer, + ape, + indices, + seq_lens, + extra_data, + kv_score_input.stride(0), + out.stride(0), + kv_score_buffer.stride(0), + kv_score_buffer.stride(1), + ape.stride(0), + HEAD_DIM=head_dim, + BLOCK_D=BLOCK_D, + BLOCK_S=64, + ) + return out + + compress_plan = _plan_as_i32(plan.compress_plan) + write_plan = _plan_as_i32(plan.write_plan) + if compress_ratio == 4: + assert extra_data is not None + assert extra_data.dim() == 2 and extra_data.shape[1] == 4 + compress_grid = (compress_plan.shape[0], num_d_chunks) + write_grid = (write_plan.shape[0], num_d_chunks) + _c4_prefill_compress_kernel[compress_grid]( + kv_score_input, + out, + kv_score_buffer, + ape, + indices, + extra_data, + compress_plan, + kv_score_input.stride(0), + out.stride(0), + kv_score_buffer.stride(0), + kv_score_buffer.stride(1), + ape.stride(0), + compress_plan.stride(0), + HEAD_DIM=head_dim, + BLOCK_D=BLOCK_D, + ) + _c4_prefill_write_kernel[write_grid]( + kv_score_input, + kv_score_buffer, + indices, + extra_data, + write_plan, + kv_score_input.stride(0), + kv_score_buffer.stride(0), + kv_score_buffer.stride(1), + write_plan.stride(0), + HEAD_DIM=head_dim, + BLOCK_D=BLOCK_D, + ) + else: + load_indices = indices if extra_data is None else extra_data + assert load_indices.dim() == 1 and load_indices.dtype == torch.int32 + compress_grid = (compress_plan.shape[0], num_d_chunks) + write_grid = (write_plan.shape[0], num_d_chunks) + _c128_prefill_compress_kernel[compress_grid]( + kv_score_input, + out, + kv_score_buffer, + ape, + load_indices, + compress_plan, + kv_score_input.stride(0), + out.stride(0), + kv_score_buffer.stride(0), + kv_score_buffer.stride(1), + ape.stride(0), + compress_plan.stride(0), + HEAD_DIM=head_dim, + BLOCK_D=BLOCK_D, + BLOCK_S=64, + ) + _c128_prefill_write_kernel[write_grid]( + kv_score_input, + kv_score_buffer, + indices, + write_plan, + kv_score_input.stride(0), + kv_score_buffer.stride(0), + kv_score_buffer.stride(1), + write_plan.stride(0), + HEAD_DIM=head_dim, + BLOCK_D=BLOCK_D, + ) + return out + + +def hip_compress_fused_norm_rope_inplace( + kv: torch.Tensor, + weight: torch.Tensor, + eps: float, + freqs_cis: torch.Tensor, + plan: Union[CompressorDecodePlan, CompressorPrefillPlan], +) -> None: + assert kv.dim() == 2 and kv.stride(-1) == 1 + assert weight.shape == (kv.shape[1],) + freqs_real = torch.view_as_real(freqs_cis).flatten(-2) + head_dim = kv.shape[1] + rope_dim = freqs_real.shape[-1] + assert head_dim >= rope_dim and rope_dim % 2 == 0 + + is_decode = _is_decode_plan(plan) + if is_decode: + handle = plan.seq_lens + else: + handle = _plan_as_i32(plan.compress_plan) + + if handle.numel() == 0: + return + + HEAD_BLOCK = triton.next_power_of_2(head_dim) + ROPE_PAIR_BLOCK = max(triton.next_power_of_2(rope_dim // 2), 1) + _compress_norm_rope_kernel[(handle.shape[0],)]( + kv, + weight, + freqs_real, + handle, + eps, + kv.stride(0), + freqs_real.stride(0), + handle.stride(0) if not is_decode else 0, + HEAD_DIM=head_dim, + ROPE_DIM=rope_dim, + HEAD_BLOCK=HEAD_BLOCK, + ROPE_PAIR_BLOCK=ROPE_PAIR_BLOCK, + COMPRESS_RATIO=plan.compress_ratio, + IS_DECODE=is_decode, + ) diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index 5284a240f..afde82d7a 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -39,6 +39,9 @@ else: FP8_MAX = torch.finfo(FP8_DTYPE).max +_arange_cache = {} + + def fp8_paged_mqa_logits_torch( q_fp8: torch.Tensor, kvcache_fp8: torch.Tensor, @@ -49,12 +52,13 @@ def fp8_paged_mqa_logits_torch( max_seq_len: int, clean_logits: bool = True, ) -> torch.Tensor: + """Vectorized implementation compatible with CUDA graph capture.""" _ = deep_gemm_metadata batch_size, _, num_heads, head_dim = q_fp8.shape block_size = kvcache_fp8.shape[1] - assert head_dim == 128, "torch reference impl hardcodes DSV4 indexer head_dim=128" - assert block_size == 64, "torch reference impl hardcodes block_size=64 cache layout" + assert head_dim == 128 + assert block_size == 64 assert q_fp8.shape == (batch_size, 1, num_heads, head_dim) assert kvcache_fp8.shape[1:] == (block_size, 1, head_dim + 4) assert weight.shape == (batch_size, num_heads) @@ -62,32 +66,85 @@ def fp8_paged_mqa_logits_torch( assert page_table.shape[0] == batch_size assert clean_logits == False - logits = page_table.new_empty((batch_size, max_seq_len), dtype=torch.float32) - for i in range(batch_size): - q = q_fp8[i, 0] - q = q.to(torch.float32) - q_scale = weight[i] - seq_len = int(seq_lens[i].item()) - assert seq_len <= max_seq_len - num_pages = (seq_len + block_size - 1) // block_size - padded_seq_len = num_pages * block_size - pages = page_table[i, :num_pages] - kvcache_fp8 = kvcache_fp8.view(-1, block_size * (head_dim + 4)) - kvcache = kvcache_fp8[pages] - SCALE_OFFSET = block_size * head_dim - kvcache_value = kvcache[..., :SCALE_OFFSET].view(dtype=FP8_DTYPE) - kvcache_scale = kvcache[..., SCALE_OFFSET:].view(dtype=torch.float32) - kvcache_value = kvcache_value.to(torch.float32) - kvcache_scale = kvcache_scale.contiguous() - kvcache_value = kvcache_value.view(padded_seq_len, head_dim) - kvcache_scale = kvcache_scale.view(padded_seq_len) - score = F.linear(kvcache_value, q) - score = F.relu(score) - score *= q_scale[None, :] - score = score.sum(dim=1) - score *= kvcache_scale - logits[i, :seq_len] = score[:seq_len] + max_num_pages = page_table.shape[1] + SCALE_OFFSET = block_size * head_dim + total_dim = block_size * (head_dim + 4) + kvcache_flat = kvcache_fp8.view(-1, total_dim) + + pages_clamped = page_table.clamp(min=0) + kvcache_gathered = kvcache_flat[pages_clamped] + + kv_values_raw = kvcache_gathered[..., :SCALE_OFFSET].contiguous() + kv_values_fp8 = kv_values_raw.view(dtype=FP8_DTYPE) + kv_values = kv_values_fp8.to(torch.float32) + kv_values = kv_values.reshape(batch_size, max_num_pages * block_size, head_dim) + + kv_scales_raw = kvcache_gathered[..., SCALE_OFFSET:].contiguous() + kv_scales = kv_scales_raw.view(dtype=torch.float32) + kv_scales = kv_scales.reshape(batch_size, max_num_pages * block_size) + + q_float = q_fp8[:, 0].to(torch.float32) + scores = torch.bmm(kv_values, q_float.transpose(1, 2)) + scores = F.relu(scores) + scores = scores * weight.unsqueeze(1) + scores = scores.sum(dim=2) + scores = scores * kv_scales + + padded_seq_len = max_num_pages * block_size + cache = _arange_cache + arange_key = f"arange_{padded_seq_len}_{scores.device}" + if arange_key not in cache: + cache[arange_key] = torch.arange(padded_seq_len, device=scores.device) + positions = cache[arange_key].unsqueeze(0) + valid_mask = positions < seq_lens.unsqueeze(1) + scores = scores.masked_fill(~valid_mask, 0.0) + + if padded_seq_len < max_seq_len: + scores = F.pad(scores, (0, max_seq_len - padded_seq_len), value=0.0) + else: + scores = scores[:, :max_seq_len] + + return scores + + +def _aiter_fp8_paged_mqa_logits( + q_fp8: torch.Tensor, + kvcache_fp8: torch.Tensor, + weight: torch.Tensor, + seq_lens: torch.Tensor, + page_table: torch.Tensor, + deep_gemm_metadata: Any, + max_seq_len: int, + clean_logits: bool = False, +) -> torch.Tensor: + """Wrapper adapting aiter's deepgemm_fp8_paged_mqa_logits to SGLang's interface.""" + from aiter.ops.triton.attention.pa_mqa_logits import ( + deepgemm_fp8_paged_mqa_logits, + ) + + batch_size = q_fp8.shape[0] + next_n = q_fp8.shape[1] + total_tokens = batch_size * next_n + _sl = seq_lens.squeeze(-1) if seq_lens.dim() == 2 else seq_lens + kv_block_size = kvcache_fp8.shape[1] + logits = torch.empty( + total_tokens, + max_seq_len, + dtype=torch.float32, + device=q_fp8.device, + ) + deepgemm_fp8_paged_mqa_logits( + q_fp8, + kvcache_fp8, + weight, + logits, + _sl.to(torch.int32), + page_table.to(torch.int32), + max_seq_len, + KVBlockSize=kv_block_size, + Preshuffle=True, + ) return logits @@ -99,6 +156,9 @@ def topk_transform_512_pytorch_vectorized( page_size: int, out_raw_indices: Optional[torch.Tensor] = None, ) -> None: + """Vectorized PyTorch fallback for topk_transform_512. + All helper tensors (arange, zeros) are cached to avoid device-tensor + creation during HIP/CUDA graph capture.""" TOPK = out_page_indices.shape[1] batch_size = scores.shape[0] @@ -108,13 +168,22 @@ def topk_transform_512_pytorch_vectorized( page_bits = (page_size - 1).bit_length() if page_size > 1 else 0 page_mask = page_size - 1 - positions = ( - torch.arange(max_seq_len, device=device).unsqueeze(0).expand(batch_size, -1) - ) + cache = _arange_cache + key_seq = f"arange_{max_seq_len}_{device}" + key_topk = f"arange_{TOPK}_{device}" + key_bs = f"arange_{batch_size}_{device}" + if key_seq not in cache: + cache[key_seq] = torch.arange(max_seq_len, device=device) + if key_topk not in cache: + cache[key_topk] = torch.arange(TOPK, device=device, dtype=torch.int32) + if key_bs not in cache: + cache[key_bs] = torch.arange(batch_size, device=device) + + positions = cache[key_seq].unsqueeze(0).expand(batch_size, -1) valid_mask = positions < seq_lens.unsqueeze(1) masked_scores = scores.clone() - masked_scores[~valid_mask] = float("-inf") + masked_scores.masked_fill_(~valid_mask, float("-inf")) actual_k = min(TOPK, max_seq_len) _, raw_indices = torch.topk( @@ -123,44 +192,28 @@ def topk_transform_512_pytorch_vectorized( raw_indices = raw_indices.to(torch.int32) if actual_k < TOPK: - padding = torch.zeros( - (batch_size, TOPK - actual_k), dtype=torch.int32, device=device - ) - raw_indices = torch.cat([raw_indices, padding], dim=1) + raw_indices = F.pad(raw_indices, (0, TOPK - actual_k), value=0) - batch_indices = ( - torch.arange(batch_size, device=device).unsqueeze(1).expand(-1, TOPK) - ) + batch_indices = cache[key_bs].unsqueeze(1).expand(-1, TOPK) gathered_scores = scores[ batch_indices.flatten(), raw_indices.clamp(min=0).flatten() ].view(batch_size, TOPK) valid_topk = gathered_scores != float("-inf") if actual_k < TOPK: - pad_mask = torch.arange(TOPK, device=device).unsqueeze(0) >= actual_k + pad_mask = cache[key_topk].unsqueeze(0) >= actual_k valid_topk = valid_topk & ~pad_mask needs_sequential = seq_lens <= TOPK - if needs_sequential.any(): - sequential_indices = ( - torch.arange(TOPK, device=device, dtype=torch.int32) - .unsqueeze(0) - .expand(batch_size, -1) - ) - sequential_valid = sequential_indices < seq_lens.unsqueeze(1) + sequential_indices = cache[key_topk].unsqueeze(0).expand(batch_size, -1) + sequential_valid = sequential_indices < seq_lens.unsqueeze(1) - raw_indices = torch.where( - needs_sequential.unsqueeze(1).expand(-1, TOPK), - torch.where( - sequential_valid, - sequential_indices, - torch.tensor(-1, device=device, dtype=torch.int32), - ), - raw_indices, - ) - valid_topk = torch.where( - needs_sequential.unsqueeze(1).expand(-1, TOPK), sequential_valid, valid_topk - ) + seq_indices_or_neg1 = sequential_indices.clone() + seq_indices_or_neg1.masked_fill_(~sequential_valid, -1) + + needs_seq_mask = needs_sequential.unsqueeze(1).expand(-1, TOPK) + raw_indices = torch.where(needs_seq_mask, seq_indices_or_neg1, raw_indices) + valid_topk = torch.where(needs_seq_mask, sequential_valid, valid_topk) page_idx = raw_indices >> page_bits offset_in_page = raw_indices & page_mask @@ -170,17 +223,13 @@ def topk_transform_512_pytorch_vectorized( page_indices = (physical_pages << page_bits) | offset_in_page page_indices = page_indices.to(torch.int32) - - page_indices = torch.where( - valid_topk, page_indices, torch.tensor(-1, device=device, dtype=torch.int32) - ) + page_indices.masked_fill_(~valid_topk, -1) out_page_indices.copy_(page_indices) if out_raw_indices is not None: - raw_indices = torch.where( - valid_topk, raw_indices, torch.tensor(-1, device=device, dtype=torch.int32) - ) + raw_indices = raw_indices.clone() + raw_indices.masked_fill_(~valid_topk, -1) out_raw_indices.copy_(raw_indices) @@ -290,18 +339,20 @@ class C4IndexerBackendMixin: positions: torch.Tensor, forward_batch: ForwardBatch, token_to_kv_pool: DeepSeekV4TokenToKVPool, + skip_compressor: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if TYPE_CHECKING: assert isinstance(self, CompressorBackendMixin) weights = c4_indexer.compute_weights(x, skip_scale=True) q_fp8, weights = c4_indexer.compute_q(q_lora, positions, weights) - self.forward_indexer_compressor( - x=x, - forward_batch=forward_batch, - layer_id=c4_indexer.layer_id, - compressor=c4_indexer.compressor, - ) + if not skip_compressor: + self.forward_indexer_compressor( + x=x, + forward_batch=forward_batch, + layer_id=c4_indexer.layer_id, + compressor=c4_indexer.compressor, + ) c4_indexer_kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer( layer_id=c4_indexer.layer_id, ) @@ -316,6 +367,7 @@ class C4IndexerBackendMixin: alt_streams: Optional[List[torch.cuda.Stream]] = None, enable_multi_stream: bool = False, q_lora_ready: Optional[torch.cuda.Event] = None, + skip_compressor: bool = False, ) -> None: if forward_batch.forward_mode.is_idle(): return @@ -354,6 +406,7 @@ class C4IndexerBackendMixin: positions=core_metadata.positions, forward_batch=forward_batch, token_to_kv_pool=token_to_kv_pool, + skip_compressor=skip_compressor, ) assert len(q_fp8.shape) == 3 @@ -372,6 +425,8 @@ class C4IndexerBackendMixin: from sglang.srt.layers.attention.dsa.tilelang_kernel import ( tilelang_fp8_paged_mqa_logits as fn, ) + elif envs.SGLANG_OPT_USE_AITER_INDEXER.get(): + fn = _aiter_fp8_paged_mqa_logits elif envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get(): fn = fp8_paged_mqa_logits_torch else: @@ -379,7 +434,8 @@ class C4IndexerBackendMixin: _c4sl = indexer_metadata.c4_seq_lens _use_tilelang = envs.SGLANG_OPT_USE_TILELANG_INDEXER.get() - if _c4sl.dim() == 1 and not _use_tilelang: + _use_aiter = envs.SGLANG_OPT_USE_AITER_INDEXER.get() + if _c4sl.dim() == 1 and not _use_tilelang and not _use_aiter: _c4sl = _c4sl.unsqueeze(-1) logits = fn( q_fp8, @@ -545,6 +601,7 @@ class C4Indexer(nn.Module): attn_backend: AttentionBackend, enable_multi_stream: bool = False, q_lora_ready: Optional[torch.cuda.Event] = None, + skip_compressor: bool = False, ) -> None: return attn_backend.forward_c4_indexer( x=x, @@ -554,4 +611,5 @@ class C4Indexer(nn.Module): alt_streams=self.alt_streams, enable_multi_stream=enable_multi_stream, q_lora_ready=q_lora_ready, + skip_compressor=skip_compressor, ) diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index c3e804103..c4d466808 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -103,7 +103,10 @@ class PagedIndexerMetadata: topk_metadata: torch.Tensor = field(init=False, repr=False) def __post_init__(self): - if envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get(): + if ( + envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get() + or envs.SGLANG_OPT_USE_AITER_INDEXER.get() + ): self.deep_gemm_metadata = None else: import deep_gemm @@ -148,14 +151,17 @@ class PagedIndexerMetadata: def copy_(self, other: "PagedIndexerMetadata"): if is_hip(): copy_fields = ["page_table", "c4_seq_lens"] + assign_fields = ["deep_gemm_metadata"] else: copy_fields = ["page_table", "c4_seq_lens", "deep_gemm_metadata"] + assign_fields = [] copy_fields += ["topk_metadata"] copy_metadata( src=other, dst=self, check_eq_fields=["page_size"], copy_fields=copy_fields, + assign_fields=assign_fields, ) diff --git a/python/sglang/srt/layers/attention/hip_flash_mla.py b/python/sglang/srt/layers/attention/hip_flash_mla.py index c22d4f38f..ae6da641f 100644 --- a/python/sglang/srt/layers/attention/hip_flash_mla.py +++ b/python/sglang/srt/layers/attention/hip_flash_mla.py @@ -12,10 +12,6 @@ def flash_mla_with_kvcache_entrypoint(backend: str, **kwargs): if is_hip(): import os - from sglang.srt.layers.attention.dsa.tilelang_kernel import ( - dpsk_v4_fp8_attention_fwd, - ) - backend = os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "tilelang") else: import flash_mla @@ -36,8 +32,19 @@ def flash_mla_with_kvcache_entrypoint(backend: str, **kwargs): return flash_mla_with_kvcache_torch(**kwargs) if backend == "tilelang": + from sglang.srt.layers.attention.dsa.tilelang_kernel import ( + dpsk_v4_fp8_attention_fwd, + ) + return dpsk_v4_fp8_attention_fwd(**kwargs) + if backend == "triton": + from sglang.srt.layers.attention.nsa.triton_decode import ( + triton_fp8_attention_fwd, + ) + + return triton_fp8_attention_fwd(**kwargs) + if backend == "kernel": return flash_mla.flash_mla_with_kvcache(**kwargs) diff --git a/python/sglang/srt/layers/attention/nsa/triton_decode/__init__.py b/python/sglang/srt/layers/attention/nsa/triton_decode/__init__.py new file mode 100644 index 000000000..7762b8bd2 --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/triton_decode/__init__.py @@ -0,0 +1,98 @@ +""" +Triton-based sparse attention decode kernels for DeepSeek V4. + +This package provides an alternative to the tilelang implementation, +controlled by the environment variable SGLANG_HACK_FLASHMLA_BACKEND=triton. +""" + +from typing import Optional, Tuple + +import torch + +from sglang.srt.layers.attention.nsa.triton_decode.triton_mla_kernels_decode_optimized import ( + triton_sparse_attn_decode, +) + + +class _KVScopeAdapter: + """Lightweight adapter providing the kv_scope interface expected by + ``triton_sparse_attn_decode``. + + The Triton kernels access four fields: + * ``blocked_k_quantized`` – the raw FP8 KV cache tensor. + * ``blocked_k`` – only ``blocked_k.shape[1]`` (block size) + is read, so we reuse the same tensor. + * ``indices_in_kvcache`` – sparse top-k page indices. + * ``topk_length`` – valid length per batch element. + """ + + __slots__ = [ + "blocked_k", + "blocked_k_quantized", + "indices_in_kvcache", + "topk_length", + ] + + def __init__( + self, + k_cache: torch.Tensor, + indices: torch.Tensor, + topk_length: Optional[torch.Tensor], + ): + self.blocked_k_quantized = k_cache + self.blocked_k = k_cache + self.indices_in_kvcache = indices + self.topk_length = topk_length + + +def triton_fp8_attention_fwd( + q: torch.Tensor, + k_cache: torch.Tensor, + head_dim_v: int, + softmax_scale: float, + indices: torch.Tensor, + attn_sink: Optional[torch.Tensor] = None, + extra_k_cache: Optional[torch.Tensor] = None, + extra_indices_in_kvcache: Optional[torch.Tensor] = None, + topk_length: Optional[torch.Tensor] = None, + extra_topk_length: Optional[torch.Tensor] = None, + **_unused, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Sparse MLA decode via Triton kernels. + + Accepts the same ``**kwargs`` dict that the caller builds for + ``flash_mla_with_kvcache`` / ``dpsk_v4_fp8_attention_fwd``, but only + uses the subset of arguments relevant to the Triton implementation. + Unused keys (``block_table``, ``cache_seqlens``, + ``tile_scheduler_metadata``, ``num_splits``, ``causal``, + ``is_fp8_kvcache``) are silently ignored via ``**_unused``. + + Returns: + ``(output, lse)`` where *output* has shape + ``[batch, seq_len, num_heads, head_dim_v]`` and *lse* has shape + ``[batch, seq_len, num_heads]``. + """ + kv_scope = _KVScopeAdapter(k_cache, indices, topk_length) + + extra_kv_scope = None + if extra_k_cache is not None: + extra_kv_scope = _KVScopeAdapter( + extra_k_cache, + extra_indices_in_kvcache, + extra_topk_length, + ) + + output, lse = triton_sparse_attn_decode( + q=q, + kv_scope=kv_scope, + extra_kv_scope=extra_kv_scope, + sm_scale=softmax_scale, + d_v=head_dim_v, + attn_sink=attn_sink, + ) + + # Triton kernel returns lse as (b, h_q, s_q); transpose to + # (b, s_q, h_q) to match the tilelang / flash_mla convention. + lse = lse.transpose(1, 2) + + return output, lse 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 new file mode 100644 index 000000000..ac9c1cced --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_common.py @@ -0,0 +1,585 @@ +""" +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 new file mode 100644 index 000000000..429a6e02f --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_dsv4.py @@ -0,0 +1,1355 @@ +""" +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 new file mode 100644 index 000000000..6167f58bf --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_fused.py @@ -0,0 +1,3089 @@ +""" +Fused Gather+Dequant+Attention Kernel for DSV4 (d_qk=512) + +This module implements a fused kernel that combines: +1. Gather: Load KV from sparse indices +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 + +import torch +import triton +import triton.language as tl + +from .triton_mla_kernels_decode_common import _bucket_total_tokens + +# ============================================================================ +# Constants for DSV4 layout +# ============================================================================ +DSV4_D_QK = 512 +DSV4_D_NOPE = 448 +DSV4_D_ROPE = 64 +DSV4_D_V = 512 +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 + + +# ============================================================================ +# 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 + kv_block_base, + nope_rope_offset, + scale_base_offset, + valid, + valid_2d, + # Query tiles + q_0, + q_1, + q_2, + q_3, + q_4, + q_5, + q_6, + q_7, + # Accumulators (passed by reference via return) + acc_0, + acc_1, + acc_2, + acc_3, + acc_4, + acc_5, + acc_6, + acc_7, + # Softmax state + m_i, + l_i, + # Other parameters + offs_tile, + sm_scale, + # Constants + TILE_SIZE: tl.constexpr, + D_NOPE: tl.constexpr, + LOG2E: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """ + Process one block of KV tokens with batch loading. + Key optimization: Load all KV tiles first, then process them. + """ + NEG_INF = float("-inf") + + scale_ptrs = kv_block_base + scale_base_offset + scale_uint8_0 = tl.load(scale_ptrs, mask=valid, other=127).to(tl.uint8) + scale_uint8_1 = tl.load(scale_ptrs + 1, mask=valid, other=127).to(tl.uint8) + scale_uint8_2 = tl.load(scale_ptrs + 2, mask=valid, other=127).to(tl.uint8) + scale_uint8_3 = tl.load(scale_ptrs + 3, mask=valid, other=127).to(tl.uint8) + scale_uint8_4 = tl.load(scale_ptrs + 4, mask=valid, other=127).to(tl.uint8) + scale_uint8_5 = tl.load(scale_ptrs + 5, mask=valid, other=127).to(tl.uint8) + scale_uint8_6 = tl.load(scale_ptrs + 6, mask=valid, other=127).to(tl.uint8) + + tile_base = kv_block_base[:, None] + nope_rope_offset[:, None] + + # Batch load all tiles + nope_uint8_0 = tl.load(tile_base + offs_tile[None, :], mask=valid_2d, other=0) + nope_uint8_1 = tl.load( + tile_base + TILE_SIZE + offs_tile[None, :], mask=valid_2d, other=0 + ) + nope_uint8_2 = tl.load( + tile_base + 2 * TILE_SIZE + offs_tile[None, :], mask=valid_2d, other=0 + ) + nope_uint8_3 = tl.load( + tile_base + 3 * TILE_SIZE + offs_tile[None, :], mask=valid_2d, other=0 + ) + nope_uint8_4 = tl.load( + tile_base + 4 * TILE_SIZE + offs_tile[None, :], mask=valid_2d, other=0 + ) + nope_uint8_5 = tl.load( + tile_base + 5 * TILE_SIZE + offs_tile[None, :], mask=valid_2d, other=0 + ) + nope_uint8_6 = tl.load( + tile_base + 6 * TILE_SIZE + offs_tile[None, :], mask=valid_2d, other=0 + ) + rope_ptrs = tile_base + D_NOPE + offs_tile[None, :] * 2 + rope_lo = tl.load(rope_ptrs, mask=valid_2d, other=0).to(tl.uint16) + rope_hi = tl.load(rope_ptrs + 1, mask=valid_2d, other=0).to(tl.uint16) + + 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) + + qk = tl.zeros([BLOCK_H, BLOCK_N], dtype=tl.float32) + + nope_fp8_0 = nope_uint8_0.to(tl.float8e4nv, bitcast=True) + kv_0 = (nope_fp8_0.to(tl.bfloat16) * scale_bf16_0[:, None]).to(tl.bfloat16) + kv_0 = tl.where(valid_2d, kv_0, 0.0) + qk += tl.dot(q_0, tl.trans(kv_0)).to(tl.float32) + + nope_fp8_1 = nope_uint8_1.to(tl.float8e4nv, bitcast=True) + kv_1 = (nope_fp8_1.to(tl.bfloat16) * scale_bf16_1[:, None]).to(tl.bfloat16) + kv_1 = tl.where(valid_2d, kv_1, 0.0) + qk += tl.dot(q_1, tl.trans(kv_1)).to(tl.float32) + + nope_fp8_2 = nope_uint8_2.to(tl.float8e4nv, bitcast=True) + kv_2 = (nope_fp8_2.to(tl.bfloat16) * scale_bf16_2[:, None]).to(tl.bfloat16) + kv_2 = tl.where(valid_2d, kv_2, 0.0) + qk += tl.dot(q_2, tl.trans(kv_2)).to(tl.float32) + + nope_fp8_3 = nope_uint8_3.to(tl.float8e4nv, bitcast=True) + kv_3 = (nope_fp8_3.to(tl.bfloat16) * scale_bf16_3[:, None]).to(tl.bfloat16) + kv_3 = tl.where(valid_2d, kv_3, 0.0) + qk += tl.dot(q_3, tl.trans(kv_3)).to(tl.float32) + + nope_fp8_4 = nope_uint8_4.to(tl.float8e4nv, bitcast=True) + kv_4 = (nope_fp8_4.to(tl.bfloat16) * scale_bf16_4[:, None]).to(tl.bfloat16) + kv_4 = tl.where(valid_2d, kv_4, 0.0) + qk += tl.dot(q_4, tl.trans(kv_4)).to(tl.float32) + + nope_fp8_5 = nope_uint8_5.to(tl.float8e4nv, bitcast=True) + kv_5 = (nope_fp8_5.to(tl.bfloat16) * scale_bf16_5[:, None]).to(tl.bfloat16) + kv_5 = tl.where(valid_2d, kv_5, 0.0) + qk += tl.dot(q_5, tl.trans(kv_5)).to(tl.float32) + + nope_fp8_6 = nope_uint8_6.to(tl.float8e4nv, bitcast=True) + kv_6 = (nope_fp8_6.to(tl.bfloat16) * scale_bf16_6[:, None]).to(tl.bfloat16) + kv_6 = tl.where(valid_2d, kv_6, 0.0) + qk += tl.dot(q_6, tl.trans(kv_6)).to(tl.float32) + + kv_7 = (rope_lo | (rope_hi << 8)).to(tl.bfloat16, bitcast=True) + kv_7 = tl.where(valid_2d, kv_7, 0.0) + qk += tl.dot(q_7, tl.trans(kv_7)).to(tl.float32) + + 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) + + acc_0 = acc_0 * alpha[:, None] + tl.dot(p_bf16, kv_0).to(tl.float32) + acc_1 = acc_1 * alpha[:, None] + tl.dot(p_bf16, kv_1).to(tl.float32) + acc_2 = acc_2 * alpha[:, None] + tl.dot(p_bf16, kv_2).to(tl.float32) + acc_3 = acc_3 * alpha[:, None] + tl.dot(p_bf16, kv_3).to(tl.float32) + acc_4 = acc_4 * alpha[:, None] + tl.dot(p_bf16, kv_4).to(tl.float32) + acc_5 = acc_5 * alpha[:, None] + tl.dot(p_bf16, kv_5).to(tl.float32) + acc_6 = acc_6 * alpha[:, None] + tl.dot(p_bf16, kv_6).to(tl.float32) + acc_7 = acc_7 * alpha[:, None] + tl.dot(p_bf16, kv_7).to(tl.float32) + + return acc_0, acc_1, acc_2, acc_3, acc_4, acc_5, acc_6, acc_7, m_new, l_new + + +# ============================================================================ +# 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": 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), + ], + key=["total_tokens_bucket", "h_q", "topk"], +) +@triton.jit +def _fused_gather_attn_dsv4_kernel( + Q, + KV_Cache, + Indices, + TopkLength, + AttnSink, + Output, + LSE, + sm_scale, + total_tokens, + total_tokens_bucket, + h_q, + topk, + num_blocks, + block_size, + s_q, + stride_q_t, + stride_q_h, + stride_q_d, + stride_kv_block, + stride_idx_t, + stride_idx_k, + stride_o_t, + stride_o_h, + stride_o_d, + stride_lse_t, + stride_lse_h, + HAS_TOPK_LENGTH: tl.constexpr, + HAS_ATTN_SINK: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """Fused gather+dequant+attention kernel for DSV4.""" + LOG2E: tl.constexpr = 1.4426950408889634 + D_NOPE: tl.constexpr = 448 + D_ROPE: tl.constexpr = 64 + TILE_SIZE: tl.constexpr = 64 + BYTES_PER_TOKEN_DATA: tl.constexpr = 576 + BYTES_PER_TOKEN_SCALE: tl.constexpr = 8 + + # OPTIMIZED: Swapped grid - pid_h first for better cache locality + pid_h = tl.program_id(0) + pid_t = tl.program_id(1) + 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 + + 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, TILE_SIZE], dtype=tl.float32) + acc_1 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_2 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_3 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_4 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_5 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_6 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_7 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + + stride_q_t_64 = tl.cast(stride_q_t, tl.int64) + q_base = Q + pid_t_64 * stride_q_t_64 + + batch_idx = pid_t // s_q + offs_tile = tl.arange(0, TILE_SIZE) + + q_0 = tl.load( + q_base + offs_h[:, None] * stride_q_h + offs_tile[None, :] * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_1 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_2 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (2 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_3 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (3 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_4 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (4 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_5 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (5 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_6 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (6 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_7 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (7 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + + # Early-exit: pre-load topk_len and skip invalid blocks + if HAS_TOPK_LENGTH: + topk_len = tl.load(TopkLength + batch_idx) + + for n_start in range(0, topk, BLOCK_N): + # Skip entire block if beyond valid topk range + should_compute = not HAS_TOPK_LENGTH or n_start < topk_len + if should_compute: + offs_n = n_start + tl.arange(0, BLOCK_N) + mask_n = offs_n < topk + + idx_ptrs = Indices + pid_t * stride_idx_t + offs_n * stride_idx_k + indices = tl.load(idx_ptrs, mask=mask_n, other=-1) + + is_invalid = indices == -1 + if HAS_TOPK_LENGTH: + is_invalid = is_invalid | (offs_n >= topk_len) + + valid = mask_n & ~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) + + stride_kv_block_64 = tl.cast(stride_kv_block, tl.int64) + kv_block_base = KV_Cache + block_idx_64 * stride_kv_block_64 + 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 + ) + + valid_2d = valid[:, None] + + # Use helper function for KV processing + acc_0, acc_1, acc_2, acc_3, acc_4, acc_5, acc_6, acc_7, m_i, l_i = ( + _process_kv_block_aggressive( + kv_block_base, + nope_rope_offset, + scale_base_offset, + valid, + valid_2d, + q_0, + q_1, + q_2, + q_3, + q_4, + q_5, + q_6, + q_7, + acc_0, + acc_1, + acc_2, + acc_3, + acc_4, + acc_5, + acc_6, + acc_7, + m_i, + l_i, + offs_tile, + sm_scale, + TILE_SIZE, + D_NOPE, + LOG2E, + BLOCK_H, + BLOCK_N, + ) + ) + + # Finalize + 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) + + acc_0 = tl.where(is_lonely_q[:, None], 0.0, acc_0 * output_scale[:, None]) + acc_1 = tl.where(is_lonely_q[:, None], 0.0, acc_1 * output_scale[:, None]) + acc_2 = tl.where(is_lonely_q[:, None], 0.0, acc_2 * output_scale[:, None]) + acc_3 = tl.where(is_lonely_q[:, None], 0.0, acc_3 * output_scale[:, None]) + acc_4 = tl.where(is_lonely_q[:, None], 0.0, acc_4 * output_scale[:, None]) + acc_5 = tl.where(is_lonely_q[:, None], 0.0, acc_5 * output_scale[:, None]) + acc_6 = tl.where(is_lonely_q[:, None], 0.0, acc_6 * output_scale[:, None]) + acc_7 = tl.where(is_lonely_q[:, None], 0.0, acc_7 * output_scale[:, None]) + lse = tl.where(is_lonely_q, float("+inf"), lse) + + stride_o_t_64 = tl.cast(stride_o_t, tl.int64) + o_base = Output + pid_t_64 * stride_o_t_64 + + # Optimized output stores with pre-computed row base pointers + # Convert to bfloat16 first (batch conversion) + o_0 = acc_0.to(tl.bfloat16) + o_1 = acc_1.to(tl.bfloat16) + o_2 = acc_2.to(tl.bfloat16) + o_3 = acc_3.to(tl.bfloat16) + o_4 = acc_4.to(tl.bfloat16) + o_5 = acc_5.to(tl.bfloat16) + o_6 = acc_6.to(tl.bfloat16) + o_7 = acc_7.to(tl.bfloat16) + + # Pre-compute row base pointers (shared across all 8 stores) + row_ptrs = o_base + offs_h[:, None] * stride_o_h + + # Store all 8 tiles with optimized pointer arithmetic + tl.store(row_ptrs + offs_tile[None, :] * stride_o_d, o_0, mask=mask_h[:, None]) + tl.store( + row_ptrs + (TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_1, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (2 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_2, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (3 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_3, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (4 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_4, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (5 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_5, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (6 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_6, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (7 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_7, + mask=mask_h[:, None], + ) + + lse_ptrs = LSE + pid_t * stride_lse_t + offs_h * stride_lse_h + 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) +BUFFER_OPS_DISABLE_THRESHOLD = 2 * 1024 * 1024 * 1024 # 2GB + + +def fused_gather_attn_decode_dsv4( + q: torch.Tensor, + kv_cache: torch.Tensor, + indices: torch.Tensor, + block_size: int, + sm_scale: float, + topk_length: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, + s_q: int = 1, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Fused gather+dequant+attention for DSV4. + Uses Split-K optimization for large topk (>= 8192). + + Args: + q: Query tensor [total_tokens, h_q, d_qk] + kv_cache: Quantized KV cache + indices: KV indices [total_tokens, topk] + block_size: Block size for KV cache + sm_scale: Softmax scale + topk_length: Optional per-batch topk length [b] + attn_sink: Optional attention sink values [h_q] + s_q: Sequence length per batch + + Returns: + output: Attention output [total_tokens, h_q, d_v] + lse: Log-sum-exp values [total_tokens, h_q] + """ + total_tokens, h_q, d_qk = q.shape + topk = indices.shape[1] + d_v = DSV4_D_V + device = q.device + + kv_uint8 = kv_cache.view(torch.uint8) + num_blocks = kv_cache.shape[0] + stride_kv_block = kv_uint8.stride(0) + kv_flat = kv_uint8.reshape(num_blocks, -1) + + if q.dtype != torch.bfloat16 or not q.is_contiguous(): + q = q.to(torch.bfloat16).contiguous() + + if not indices.is_contiguous(): + indices = indices.contiguous() + + kv_cache_size = stride_kv_block * num_blocks + disable_buffer_ops = kv_cache_size > BUFFER_OPS_DISABLE_THRESHOLD + + # Use Split-K for large topk + if topk >= SPLITK_TOPK_THRESHOLD: + split_k = _select_split_k(topk, h_q, total_tokens) + topk_per_split = (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 + ) + 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) + + topk_length_tensor = topk_length if topk_length is not None else lse[:1, 0] + attn_sink_tensor = attn_sink if attn_sink is not None else lse[0, :] + + # Use autotuned grid + grid_splitk = lambda meta: ( + triton.cdiv(h_q, meta["BLOCK_H"]), + total_tokens, + split_k, + ) + + def run_splitk_kernel(): + _fused_gather_attn_dsv4_splitk_kernel[grid_splitk]( + q, + kv_flat, + indices, + topk_length_tensor, + partial_output, + partial_lse, + sm_scale, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + topk, + num_blocks, + block_size, + s_q, + topk_per_split, + q.stride(0), + q.stride(1), + q.stride(2), + stride_kv_block, + indices.stride(0), + indices.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), + HAS_TOPK_LENGTH=topk_length is not None, + ) + + if disable_buffer_ops: + with triton.knobs.amd.scope(): + triton.knobs.amd.use_buffer_ops = False + run_splitk_kernel() + else: + run_splitk_kernel() + + # Use autotuned combine kernel for split_k=8 + if split_k == 8: + # Autotuned kernel - grid is determined by autotune + grid_combine = lambda meta: ( + total_tokens, + triton.cdiv(h_q, meta["BLOCK_H"]), + ) + _combine_splitk_kernel_8_optimized[grid_combine]( + partial_output, + partial_lse, + attn_sink_tensor, + output, + lse, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + d_v, + 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), + output.stride(0), + output.stride(1), + output.stride(2), + lse.stride(0), + lse.stride(1), + HAS_ATTN_SINK=attn_sink is not None, + ) + else: + BLOCK_H_COMBINE = 16 + BLOCK_D_COMBINE = 128 + grid_combine = (total_tokens, triton.cdiv(h_q, BLOCK_H_COMBINE)) + + # Select appropriate combine kernel based on split_k + if split_k == 2: + combine_kernel = _combine_splitk_kernel_2 + elif split_k == 4: + combine_kernel = _combine_splitk_kernel + else: + raise ValueError(f"Unsupported split_k: {split_k}") + + combine_kernel[grid_combine]( + partial_output, + partial_lse, + attn_sink_tensor, + output, + lse, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + d_v, + 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), + output.stride(0), + output.stride(1), + output.stride(2), + lse.stride(0), + lse.stride(1), + HAS_ATTN_SINK=attn_sink is not None, + BLOCK_H=BLOCK_H_COMBINE, + BLOCK_D=BLOCK_D_COMBINE, + num_warps=4, + num_stages=1, + ) + + return output, lse + + # Use original kernel for smaller topk + 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) + + topk_length_tensor = topk_length if topk_length is not None else lse[:1, 0] + attn_sink_tensor = attn_sink if attn_sink is not None else lse[0, :] + + grid = lambda meta: (triton.cdiv(h_q, meta["BLOCK_H"]), total_tokens) + + def run_kernel(): + _fused_gather_attn_dsv4_kernel[grid]( + q, + kv_flat, + indices, + topk_length_tensor, + attn_sink_tensor, + output, + lse, + sm_scale, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + topk, + num_blocks, + block_size, + s_q, + q.stride(0), + q.stride(1), + q.stride(2), + stride_kv_block, + indices.stride(0), + indices.stride(1), + output.stride(0), + output.stride(1), + output.stride(2), + lse.stride(0), + lse.stride(1), + HAS_TOPK_LENGTH=topk_length is not None, + HAS_ATTN_SINK=attn_sink is not None, + ) + + if disable_buffer_ops: + with triton.knobs.amd.scope(): + triton.knobs.amd.use_buffer_ops = False + run_kernel() + else: + run_kernel() + + 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. + + 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) + """ + h_q = named_args.get("h_q", 128) + 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), + 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": 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), + # warps=8: for memory-bound scenarios + triton.Config({"BLOCK_H": 16, "BLOCK_N": 64}, num_warps=8, num_stages=1), + triton.Config({"BLOCK_H": 64, "BLOCK_N": 64}, num_warps=8, num_stages=1), + ], + key=["total_tokens_bucket", "h_q", "topk_main", "topk_extra"], + prune_configs_by={"early_config_prune": _prune_dual_scope_configs}, +) +@triton.jit +def _fused_gather_attn_dsv4_dual_scope_kernel( + Q, + KV_Cache_Main, + Indices_Main, + TopkLength_Main, + KV_Cache_Extra, + Indices_Extra, + TopkLength_Extra, + AttnSink, + Output, + LSE, + sm_scale, + total_tokens, + total_tokens_bucket, + h_q, + topk_main, + num_blocks_main, + block_size_main, + topk_extra, + num_blocks_extra, + block_size_extra, + s_q, + stride_q_t, + stride_q_h, + stride_q_d, + stride_kv_block_main, + stride_kv_block_extra, + stride_idx_main_t, + stride_idx_main_k, + stride_idx_extra_t, + stride_idx_extra_k, + stride_o_t, + stride_o_h, + stride_o_d, + stride_lse_t, + stride_lse_h, + HAS_TOPK_LENGTH_MAIN: tl.constexpr, + HAS_TOPK_LENGTH_EXTRA: tl.constexpr, + HAS_ATTN_SINK: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """ + OPTIMIZED fused gather+dequant+attention kernel for DSV4 with dual scope. + + This version uses a helper function (_process_kv_block_aggressive) to + eliminate the ~200 lines of duplicated code between MAIN and EXTRA scope + processing loops. + + The kernel processes: + 1. MAIN scope: topk_main tokens from KV_Cache_Main + 2. EXTRA scope: topk_extra tokens from KV_Cache_Extra + + Both scopes contribute to the same online softmax accumulator. + """ + LOG2E: tl.constexpr = 1.4426950408889634 + D_NOPE: tl.constexpr = 448 + D_ROPE: tl.constexpr = 64 + TILE_SIZE: tl.constexpr = 64 + BYTES_PER_TOKEN_DATA: tl.constexpr = 576 + BYTES_PER_TOKEN_SCALE: tl.constexpr = 8 + + # OPTIMIZED: Swapped grid - pid_h first for better cache locality + pid_h = tl.program_id(0) + pid_t = tl.program_id(1) + 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 + + # Initialize accumulators + 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, TILE_SIZE], dtype=tl.float32) + acc_1 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_2 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_3 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_4 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_5 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_6 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_7 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + + stride_q_t_64 = tl.cast(stride_q_t, tl.int64) + q_base = Q + pid_t_64 * stride_q_t_64 + + batch_idx = pid_t // s_q + offs_tile = tl.arange(0, TILE_SIZE) + + # Load Q tiles (shared by both scopes) + q_0 = tl.load( + q_base + offs_h[:, None] * stride_q_h + offs_tile[None, :] * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_1 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_2 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (2 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_3 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (3 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_4 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (4 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_5 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (5 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_6 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (6 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_7 = tl.load( + q_base + + offs_h[:, None] * stride_q_h + + (7 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + + # ======================================================================== + # Process MAIN scope + # ======================================================================== + # Early-exit: pre-load topk_len and skip invalid blocks + if HAS_TOPK_LENGTH_MAIN: + topk_len = tl.load(TopkLength_Main + batch_idx) + + for n_start in range(0, topk_main, BLOCK_N): + # Skip entire block if beyond valid topk range + should_compute = not HAS_TOPK_LENGTH_MAIN or n_start < topk_len + if should_compute: + offs_n = n_start + tl.arange(0, BLOCK_N) + mask_n = offs_n < topk_main + + idx_ptrs = ( + Indices_Main + pid_t * stride_idx_main_t + offs_n * stride_idx_main_k + ) + indices = tl.load(idx_ptrs, mask=mask_n, other=-1) + + is_invalid = indices == -1 + if HAS_TOPK_LENGTH_MAIN: + is_invalid = is_invalid | (offs_n >= topk_len) + + valid = mask_n & ~is_invalid + indices_clamped = tl.maximum(indices, 0) + + block_idx = indices_clamped // block_size_main + offset_in_block = indices_clamped % block_size_main + + block_idx_64 = block_idx.to(tl.int64) + offset_in_block_64 = offset_in_block.to(tl.int64) + + stride_kv_block_main_64 = tl.cast(stride_kv_block_main, tl.int64) + kv_block_base = KV_Cache_Main + block_idx_64 * stride_kv_block_main_64 + nope_rope_offset = offset_in_block_64 * BYTES_PER_TOKEN_DATA + scale_base_offset = ( + block_size_main * BYTES_PER_TOKEN_DATA + + offset_in_block_64 * BYTES_PER_TOKEN_SCALE + ) + + valid_2d = valid[:, None] + + # Use helper function for KV processing + acc_0, acc_1, acc_2, acc_3, acc_4, acc_5, acc_6, acc_7, m_i, l_i = ( + _process_kv_block_aggressive( + kv_block_base, + nope_rope_offset, + scale_base_offset, + valid, + valid_2d, + q_0, + q_1, + q_2, + q_3, + q_4, + q_5, + q_6, + q_7, + acc_0, + acc_1, + acc_2, + acc_3, + acc_4, + acc_5, + acc_6, + acc_7, + m_i, + l_i, + offs_tile, + sm_scale, + TILE_SIZE, + D_NOPE, + LOG2E, + BLOCK_H, + BLOCK_N, + ) + ) + + # ======================================================================== + # Process EXTRA scope + # ======================================================================== + # Early-exit: pre-load topk_len and skip invalid blocks + if HAS_TOPK_LENGTH_EXTRA: + topk_len = tl.load(TopkLength_Extra + batch_idx) + + for n_start in range(0, topk_extra, BLOCK_N): + # Skip entire block if beyond valid topk range + should_compute = not HAS_TOPK_LENGTH_EXTRA or n_start < topk_len + if should_compute: + offs_n = n_start + tl.arange(0, BLOCK_N) + mask_n = offs_n < topk_extra + + idx_ptrs = ( + Indices_Extra + pid_t * stride_idx_extra_t + offs_n * stride_idx_extra_k + ) + indices = tl.load(idx_ptrs, mask=mask_n, other=-1) + + is_invalid = indices == -1 + if HAS_TOPK_LENGTH_EXTRA: + is_invalid = is_invalid | (offs_n >= topk_len) + + valid = mask_n & ~is_invalid + indices_clamped = tl.maximum(indices, 0) + + block_idx = indices_clamped // block_size_extra + offset_in_block = indices_clamped % block_size_extra + + block_idx_64 = block_idx.to(tl.int64) + offset_in_block_64 = offset_in_block.to(tl.int64) + + stride_kv_block_extra_64 = tl.cast(stride_kv_block_extra, tl.int64) + kv_block_base = KV_Cache_Extra + block_idx_64 * stride_kv_block_extra_64 + nope_rope_offset = offset_in_block_64 * BYTES_PER_TOKEN_DATA + scale_base_offset = ( + block_size_extra * BYTES_PER_TOKEN_DATA + + offset_in_block_64 * BYTES_PER_TOKEN_SCALE + ) + + valid_2d = valid[:, None] + + # Use helper function for KV processing + acc_0, acc_1, acc_2, acc_3, acc_4, acc_5, acc_6, acc_7, m_i, l_i = ( + _process_kv_block_aggressive( + kv_block_base, + nope_rope_offset, + scale_base_offset, + valid, + valid_2d, + q_0, + q_1, + q_2, + q_3, + q_4, + q_5, + q_6, + q_7, + acc_0, + acc_1, + acc_2, + acc_3, + acc_4, + acc_5, + acc_6, + acc_7, + m_i, + l_i, + offs_tile, + sm_scale, + TILE_SIZE, + D_NOPE, + LOG2E, + BLOCK_H, + BLOCK_N, + ) + ) + + # ======================================================================== + # Finalize: compute LSE and output + # ======================================================================== + lse = m_i + tl.math.log2(tl.where(l_i == 0.0, 1.0, l_i)) / LOG2E + is_lonely_q = l_i == 0.0 + + # Compute output scale + 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) + + # Apply output scaling and handle lonely queries + acc_0 = tl.where(is_lonely_q[:, None], 0.0, acc_0 * output_scale[:, None]) + acc_1 = tl.where(is_lonely_q[:, None], 0.0, acc_1 * output_scale[:, None]) + acc_2 = tl.where(is_lonely_q[:, None], 0.0, acc_2 * output_scale[:, None]) + acc_3 = tl.where(is_lonely_q[:, None], 0.0, acc_3 * output_scale[:, None]) + acc_4 = tl.where(is_lonely_q[:, None], 0.0, acc_4 * output_scale[:, None]) + acc_5 = tl.where(is_lonely_q[:, None], 0.0, acc_5 * output_scale[:, None]) + acc_6 = tl.where(is_lonely_q[:, None], 0.0, acc_6 * output_scale[:, None]) + acc_7 = tl.where(is_lonely_q[:, None], 0.0, acc_7 * output_scale[:, None]) + lse = tl.where(is_lonely_q, float("+inf"), lse) + + stride_o_t_64 = tl.cast(stride_o_t, tl.int64) + o_base = Output + pid_t_64 * stride_o_t_64 + + # Optimized output stores with pre-computed row base pointers + # Convert to bfloat16 first (batch conversion) + o_0 = acc_0.to(tl.bfloat16) + o_1 = acc_1.to(tl.bfloat16) + o_2 = acc_2.to(tl.bfloat16) + o_3 = acc_3.to(tl.bfloat16) + o_4 = acc_4.to(tl.bfloat16) + o_5 = acc_5.to(tl.bfloat16) + o_6 = acc_6.to(tl.bfloat16) + o_7 = acc_7.to(tl.bfloat16) + + # Pre-compute row base pointers (shared across all 8 stores) + row_ptrs = o_base + offs_h[:, None] * stride_o_h + + # Store all 8 tiles with optimized pointer arithmetic + tl.store(row_ptrs + offs_tile[None, :] * stride_o_d, o_0, mask=mask_h[:, None]) + tl.store( + row_ptrs + (TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_1, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (2 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_2, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (3 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_3, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (4 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_4, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (5 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_5, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (6 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_6, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (7 * TILE_SIZE + offs_tile[None, :]) * stride_o_d, + o_7, + mask=mask_h[:, None], + ) + + lse_ptrs = LSE + pid_t * stride_lse_t + offs_h * stride_lse_h + 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. + 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), + ], + 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( + Q, + KV_Cache_Main, + Indices_Main, + TopkLength_Main, + KV_Cache_Extra, + Indices_Extra, + TopkLength_Extra, + PartialOutput, + PartialLSE, + sm_scale, + total_tokens, + total_tokens_bucket, + h_q, + topk_main, + num_blocks_main, + block_size_main, + topk_extra, + num_blocks_extra, + block_size_extra, + s_q, + topk_per_split, + stride_q_t, + stride_q_h, + stride_q_d, + stride_kv_block_main, + stride_kv_block_extra, + stride_idx_main_t, + stride_idx_main_k, + stride_idx_extra_t, + stride_idx_extra_k, + stride_po_s, + stride_po_t, + stride_po_h, + stride_po_d, + stride_plse_s, + stride_plse_t, + stride_plse_h, + HAS_TOPK_LENGTH_MAIN: tl.constexpr, + HAS_TOPK_LENGTH_EXTRA: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """ + Split-K fused gather+dequant+attention kernel for DSV4 with dual scope. + + This kernel processes a portion of the combined topk range (main + extra). + Each split handles topk_per_split tokens from the combined range. + """ + LOG2E: tl.constexpr = 1.4426950408889634 + D_NOPE: tl.constexpr = 448 + TILE_SIZE: tl.constexpr = 64 + BYTES_PER_TOKEN_DATA: tl.constexpr = 576 + BYTES_PER_TOKEN_SCALE: tl.constexpr = 8 + + pid_h = tl.program_id(0) + pid_t = 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 + + # Calculate the range for this split + total_topk = topk_main + topk_extra + k_start = pid_k * topk_per_split + k_end = tl.minimum(k_start + topk_per_split, total_topk) + + # Initialize accumulators + 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, TILE_SIZE], dtype=tl.float32) + acc_1 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_2 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_3 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_4 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_5 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_6 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_7 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + + stride_q_t_64 = tl.cast(stride_q_t, tl.int64) + q_base = Q + pid_t_64 * stride_q_t_64 + + batch_idx = pid_t // s_q + offs_tile = tl.arange(0, TILE_SIZE) + + # Load Q tiles (shared by both scopes) + q_row_base = q_base + offs_h[:, None] * stride_q_h + q_0 = tl.load( + q_row_base + offs_tile[None, :] * stride_q_d, mask=mask_h[:, None], other=0.0 + ).to(tl.bfloat16) + q_1 = tl.load( + q_row_base + (TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_2 = tl.load( + q_row_base + (2 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_3 = tl.load( + q_row_base + (3 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_4 = tl.load( + q_row_base + (4 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_5 = tl.load( + q_row_base + (5 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_6 = tl.load( + q_row_base + (6 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_7 = tl.load( + q_row_base + (7 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + + stride_kv_block_main_64 = tl.cast(stride_kv_block_main, tl.int64) + stride_kv_block_extra_64 = tl.cast(stride_kv_block_extra, tl.int64) + + # Process the combined range [k_start, k_end) + # First, process MAIN scope portion (indices 0 to topk_main-1) + main_start = k_start + main_end = tl.minimum(k_end, topk_main) + + # Early-exit: pre-load topk_len and skip invalid blocks + if HAS_TOPK_LENGTH_MAIN: + topk_len = tl.load(TopkLength_Main + batch_idx) + + for n_start in range(main_start, main_end, BLOCK_N): + # Skip entire block if beyond valid topk range + should_compute = not HAS_TOPK_LENGTH_MAIN or n_start < topk_len + if should_compute: + offs_n = n_start + tl.arange(0, BLOCK_N) + mask_n = offs_n < main_end + + idx_ptrs = ( + Indices_Main + pid_t * stride_idx_main_t + offs_n * stride_idx_main_k + ) + indices = tl.load(idx_ptrs, mask=mask_n, other=-1) + + is_invalid = indices == -1 + if HAS_TOPK_LENGTH_MAIN: + is_invalid = is_invalid | (offs_n >= topk_len) + + valid = mask_n & ~is_invalid + indices_clamped = tl.maximum(indices, 0) + + block_idx = indices_clamped // block_size_main + offset_in_block = indices_clamped % block_size_main + + block_idx_64 = block_idx.to(tl.int64) + offset_in_block_64 = offset_in_block.to(tl.int64) + + kv_block_base = KV_Cache_Main + block_idx_64 * stride_kv_block_main_64 + nope_rope_offset = offset_in_block_64 * BYTES_PER_TOKEN_DATA + scale_base_offset = ( + block_size_main * BYTES_PER_TOKEN_DATA + + offset_in_block_64 * BYTES_PER_TOKEN_SCALE + ) + + valid_2d = valid[:, None] + + acc_0, acc_1, acc_2, acc_3, acc_4, acc_5, acc_6, acc_7, m_i, l_i = ( + _process_kv_block_aggressive( + kv_block_base, + nope_rope_offset, + scale_base_offset, + valid, + valid_2d, + q_0, + q_1, + q_2, + q_3, + q_4, + q_5, + q_6, + q_7, + acc_0, + acc_1, + acc_2, + acc_3, + acc_4, + acc_5, + acc_6, + acc_7, + m_i, + l_i, + offs_tile, + sm_scale, + TILE_SIZE, + D_NOPE, + LOG2E, + BLOCK_H, + BLOCK_N, + ) + ) + + # Process EXTRA scope portion (indices topk_main to topk_main+topk_extra-1) + extra_global_start = tl.maximum(k_start, topk_main) + extra_global_end = k_end + + # Early-exit: pre-load topk_len and skip invalid blocks + if HAS_TOPK_LENGTH_EXTRA: + topk_len = tl.load(TopkLength_Extra + batch_idx) + + for n_global in range(extra_global_start, extra_global_end, BLOCK_N): + # Skip entire block if beyond valid topk range + should_compute = not HAS_TOPK_LENGTH_EXTRA or (n_global - topk_main) < topk_len + if should_compute: + offs_n_local = (n_global - topk_main) + tl.arange(0, BLOCK_N) + offs_n_global = n_global + tl.arange(0, BLOCK_N) + mask_n = offs_n_global < extra_global_end + + idx_ptrs = ( + Indices_Extra + + pid_t * stride_idx_extra_t + + offs_n_local * stride_idx_extra_k + ) + indices = tl.load(idx_ptrs, mask=mask_n, other=-1) + + is_invalid = indices == -1 + if HAS_TOPK_LENGTH_EXTRA: + is_invalid = is_invalid | (offs_n_local >= topk_len) + + valid = mask_n & ~is_invalid + indices_clamped = tl.maximum(indices, 0) + + block_idx = indices_clamped // block_size_extra + offset_in_block = indices_clamped % block_size_extra + + block_idx_64 = block_idx.to(tl.int64) + offset_in_block_64 = offset_in_block.to(tl.int64) + + kv_block_base = KV_Cache_Extra + block_idx_64 * stride_kv_block_extra_64 + nope_rope_offset = offset_in_block_64 * BYTES_PER_TOKEN_DATA + scale_base_offset = ( + block_size_extra * BYTES_PER_TOKEN_DATA + + offset_in_block_64 * BYTES_PER_TOKEN_SCALE + ) + + valid_2d = valid[:, None] + + acc_0, acc_1, acc_2, acc_3, acc_4, acc_5, acc_6, acc_7, m_i, l_i = ( + _process_kv_block_aggressive( + kv_block_base, + nope_rope_offset, + scale_base_offset, + valid, + valid_2d, + q_0, + q_1, + q_2, + q_3, + q_4, + q_5, + q_6, + q_7, + acc_0, + acc_1, + acc_2, + acc_3, + acc_4, + acc_5, + acc_6, + acc_7, + m_i, + l_i, + offs_tile, + sm_scale, + TILE_SIZE, + D_NOPE, + LOG2E, + BLOCK_H, + BLOCK_N, + ) + ) + + # Finalize: compute partial LSE and store partial output + lse = m_i + tl.math.log2(tl.where(l_i == 0.0, 1.0, l_i)) / LOG2E + is_lonely_q = l_i == 0.0 + + output_scale = tl.where(l_i == 0.0, 0.0, 1.0 / l_i) + + acc_0 = tl.where(is_lonely_q[:, None], 0.0, acc_0 * output_scale[:, None]) + acc_1 = tl.where(is_lonely_q[:, None], 0.0, acc_1 * output_scale[:, None]) + acc_2 = tl.where(is_lonely_q[:, None], 0.0, acc_2 * output_scale[:, None]) + acc_3 = tl.where(is_lonely_q[:, None], 0.0, acc_3 * output_scale[:, None]) + acc_4 = tl.where(is_lonely_q[:, None], 0.0, acc_4 * output_scale[:, None]) + acc_5 = tl.where(is_lonely_q[:, None], 0.0, acc_5 * output_scale[:, None]) + acc_6 = tl.where(is_lonely_q[:, None], 0.0, acc_6 * output_scale[:, None]) + acc_7 = tl.where(is_lonely_q[:, None], 0.0, acc_7 * output_scale[:, None]) + lse = tl.where(is_lonely_q, float("+inf"), lse) + + # Store partial output + 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 + + # Store partial output as float32 for better precision in combine kernel + row_ptrs = po_base + offs_h[:, None] * stride_po_h + + tl.store(row_ptrs + offs_tile[None, :] * stride_po_d, acc_0, mask=mask_h[:, None]) + tl.store( + row_ptrs + (TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_1, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (2 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_2, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (3 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_3, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (4 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_4, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (5 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_5, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (6 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_6, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (7 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_7, + mask=mask_h[:, None], + ) + + # Store partial LSE + stride_plse_s_64 = tl.cast(stride_plse_s, tl.int64) + stride_plse_t_64 = tl.cast(stride_plse_t, tl.int64) + lse_ptrs = ( + PartialLSE + + pid_k * stride_plse_s_64 + + pid_t_64 * stride_plse_t_64 + + offs_h * stride_plse_h + ) + tl.store(lse_ptrs, lse, mask=mask_h) + + +def fused_gather_attn_decode_dsv4_dual_scope( + q: torch.Tensor, + kv_cache_main: torch.Tensor, + indices_main: torch.Tensor, + block_size_main: int, + kv_cache_extra: torch.Tensor, + indices_extra: torch.Tensor, + block_size_extra: int, + sm_scale: float, + topk_length_main: Optional[torch.Tensor] = None, + topk_length_extra: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, + s_q: int = 1, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Fused gather+dequant+attention for DSV4 with dual scope (main + extra). + Uses Split-K optimization for large total_topk (>= SPLITK_TOPK_THRESHOLD). + + Args: + q: Query tensor [total_tokens, h_q, d_qk] + kv_cache_main: Quantized main KV cache + indices_main: Main KV indices [total_tokens, topk_main] + block_size_main: Block size for main KV cache + kv_cache_extra: Quantized extra KV cache + indices_extra: Extra KV indices [total_tokens, topk_extra] + block_size_extra: Block size for extra KV cache + sm_scale: Softmax scale + topk_length_main: Optional per-batch topk length for main [b] + topk_length_extra: Optional per-batch topk length for extra [b] + attn_sink: Optional attention sink values [h_q] + s_q: Sequence length per batch + + Returns: + output: Attention output [total_tokens, h_q, d_v] + lse: Log-sum-exp values [total_tokens, h_q] + """ + total_tokens, h_q, d_qk = q.shape + topk_main = indices_main.shape[1] + topk_extra = indices_extra.shape[1] + total_topk = topk_main + topk_extra + d_v = DSV4_D_V + device = q.device + + # Prepare main KV cache + kv_uint8_main = kv_cache_main.view(torch.uint8) + num_blocks_main = kv_cache_main.shape[0] + stride_kv_block_main = kv_uint8_main.stride(0) + kv_flat_main = kv_uint8_main.reshape(num_blocks_main, -1) + + # Prepare extra KV cache + kv_uint8_extra = kv_cache_extra.view(torch.uint8) + num_blocks_extra = kv_cache_extra.shape[0] + stride_kv_block_extra = kv_uint8_extra.stride(0) + kv_flat_extra = kv_uint8_extra.reshape(num_blocks_extra, -1) + + if q.dtype != torch.bfloat16 or not q.is_contiguous(): + q = q.to(torch.bfloat16).contiguous() + + if not indices_main.is_contiguous(): + indices_main = indices_main.contiguous() + if not indices_extra.is_contiguous(): + indices_extra = indices_extra.contiguous() + + kv_cache_size_main = stride_kv_block_main * num_blocks_main + kv_cache_size_extra = stride_kv_block_extra * num_blocks_extra + disable_buffer_ops = ( + kv_cache_size_main > BUFFER_OPS_DISABLE_THRESHOLD + or kv_cache_size_extra > BUFFER_OPS_DISABLE_THRESHOLD + ) + + # 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 + ) + 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_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) + 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 + ) + 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) + + topk_length_main_tensor = ( + topk_length_main if topk_length_main is not None else lse[:1, 0] + ) + topk_length_extra_tensor = ( + topk_length_extra if topk_length_extra is not None else lse[:1, 0] + ) + attn_sink_tensor = attn_sink if attn_sink is not None else lse[0, :] + + grid_splitk = lambda meta: ( + triton.cdiv(h_q, meta["BLOCK_H"]), + total_tokens, + split_k, + ) + + def run_splitk_kernel(): + _fused_gather_attn_dsv4_dual_scope_splitk_kernel[grid_splitk]( + q, + kv_flat_main, + indices_main, + topk_length_main_tensor, + kv_flat_extra, + indices_extra, + topk_length_extra_tensor, + partial_output, + partial_lse, + sm_scale, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + topk_main, + num_blocks_main, + block_size_main, + topk_extra, + num_blocks_extra, + block_size_extra, + s_q, + topk_per_split, + q.stride(0), + q.stride(1), + q.stride(2), + stride_kv_block_main, + stride_kv_block_extra, + indices_main.stride(0), + indices_main.stride(1), + indices_extra.stride(0), + indices_extra.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), + HAS_TOPK_LENGTH_MAIN=topk_length_main is not None, + HAS_TOPK_LENGTH_EXTRA=topk_length_extra is not None, + ) + + if disable_buffer_ops: + with triton.knobs.amd.scope(): + triton.knobs.amd.use_buffer_ops = False + run_splitk_kernel() + else: + run_splitk_kernel() + + # Use appropriate combine kernel based on split_k + if split_k == 8: + grid_combine = lambda meta: ( + total_tokens, + triton.cdiv(h_q, meta["BLOCK_H"]), + ) + _combine_splitk_kernel_8_optimized[grid_combine]( + partial_output, + partial_lse, + attn_sink_tensor, + output, + lse, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + d_v, + 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), + output.stride(0), + output.stride(1), + output.stride(2), + lse.stride(0), + lse.stride(1), + HAS_ATTN_SINK=attn_sink is not None, + ) + else: + BLOCK_H_COMBINE = 16 + BLOCK_D_COMBINE = 128 + grid_combine = (total_tokens, triton.cdiv(h_q, BLOCK_H_COMBINE)) + + if split_k == 2: + combine_kernel = _combine_splitk_kernel_2 + elif split_k == 4: + combine_kernel = _combine_splitk_kernel + else: + raise ValueError(f"Unsupported split_k: {split_k}") + + combine_kernel[grid_combine]( + partial_output, + partial_lse, + attn_sink_tensor, + output, + lse, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + d_v, + 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), + output.stride(0), + output.stride(1), + output.stride(2), + lse.stride(0), + lse.stride(1), + HAS_ATTN_SINK=attn_sink is not None, + BLOCK_H=BLOCK_H_COMBINE, + BLOCK_D=BLOCK_D_COMBINE, + num_warps=4, + num_stages=1, + ) + + return output, lse + + # Use original kernel for smaller total_topk + 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) + + topk_length_main_tensor = ( + topk_length_main if topk_length_main is not None else lse[:1, 0] + ) + topk_length_extra_tensor = ( + topk_length_extra if topk_length_extra is not None else lse[:1, 0] + ) + attn_sink_tensor = attn_sink if attn_sink is not None else lse[0, :] + + grid = lambda meta: (triton.cdiv(h_q, meta["BLOCK_H"]), total_tokens) + + def run_kernel(): + _fused_gather_attn_dsv4_dual_scope_kernel[grid]( + q, + kv_flat_main, + indices_main, + topk_length_main_tensor, + kv_flat_extra, + indices_extra, + topk_length_extra_tensor, + attn_sink_tensor, + output, + lse, + sm_scale, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + topk_main, + num_blocks_main, + block_size_main, + topk_extra, + num_blocks_extra, + block_size_extra, + s_q, + q.stride(0), + q.stride(1), + q.stride(2), + stride_kv_block_main, + stride_kv_block_extra, + indices_main.stride(0), + indices_main.stride(1), + indices_extra.stride(0), + indices_extra.stride(1), + output.stride(0), + output.stride(1), + output.stride(2), + lse.stride(0), + lse.stride(1), + HAS_TOPK_LENGTH_MAIN=topk_length_main is not None, + HAS_TOPK_LENGTH_EXTRA=topk_length_extra is not None, + HAS_ATTN_SINK=attn_sink is not None, + ) + + if disable_buffer_ops: + with triton.knobs.amd.scope(): + triton.knobs.amd.use_buffer_ops = False + run_kernel() + else: + run_kernel() + + return output, lse + + +# ============================================================================ +# 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": 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": 128, "BLOCK_N": 16}, num_warps=4, num_stages=1), + triton.Config({"BLOCK_H": 128, "BLOCK_N": 32}, num_warps=4, num_stages=1), + ], + key=["total_tokens_bucket", "h_q", "topk_per_split"], +) +@triton.jit +def _fused_gather_attn_dsv4_splitk_kernel( + Q, + KV_Cache, + Indices, + TopkLength, + PartialOutput, + PartialLSE, + sm_scale, + total_tokens, + total_tokens_bucket, + h_q, + topk, + num_blocks, + block_size, + s_q, + topk_per_split, + stride_q_t, + stride_q_h, + stride_q_d, + stride_kv_block, + stride_idx_t, + stride_idx_k, + stride_po_s, + stride_po_t, + stride_po_h, + stride_po_d, + stride_plse_s, + stride_plse_t, + stride_plse_h, + HAS_TOPK_LENGTH: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """Split-K fused gather+dequant+attention kernel for DSV4.""" + LOG2E: tl.constexpr = 1.4426950408889634 + D_NOPE: tl.constexpr = 448 + TILE_SIZE: tl.constexpr = 64 + BYTES_PER_TOKEN_DATA: tl.constexpr = 576 + BYTES_PER_TOKEN_SCALE: tl.constexpr = 8 + + pid_h = tl.program_id(0) + pid_t = 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 + + k_start = pid_k * topk_per_split + k_end = tl.minimum(k_start + topk_per_split, 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, TILE_SIZE], dtype=tl.float32) + acc_1 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_2 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_3 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_4 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_5 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_6 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + acc_7 = tl.zeros([BLOCK_H, TILE_SIZE], dtype=tl.float32) + + stride_q_t_64 = tl.cast(stride_q_t, tl.int64) + q_base = Q + pid_t_64 * stride_q_t_64 + + batch_idx = pid_t // s_q + offs_tile = tl.arange(0, TILE_SIZE) + + q_row_base = q_base + offs_h[:, None] * stride_q_h + q_0 = tl.load( + q_row_base + offs_tile[None, :] * stride_q_d, mask=mask_h[:, None], other=0.0 + ).to(tl.bfloat16) + q_1 = tl.load( + q_row_base + (TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_2 = tl.load( + q_row_base + (2 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_3 = tl.load( + q_row_base + (3 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_4 = tl.load( + q_row_base + (4 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_5 = tl.load( + q_row_base + (5 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_6 = tl.load( + q_row_base + (6 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + q_7 = tl.load( + q_row_base + (7 * TILE_SIZE + offs_tile[None, :]) * stride_q_d, + mask=mask_h[:, None], + other=0.0, + ).to(tl.bfloat16) + + stride_kv_block_64 = tl.cast(stride_kv_block, tl.int64) + + # Early-exit: pre-load topk_len and skip invalid blocks + if HAS_TOPK_LENGTH: + topk_len = tl.load(TopkLength + batch_idx) + + for n_start in range(k_start, k_end, BLOCK_N): + # Skip entire block if beyond valid topk range + should_compute = not HAS_TOPK_LENGTH or n_start < topk_len + if should_compute: + offs_n = n_start + tl.arange(0, BLOCK_N) + mask_n = offs_n < k_end + + idx_ptrs = Indices + pid_t * stride_idx_t + offs_n * stride_idx_k + indices = tl.load(idx_ptrs, mask=mask_n, other=-1) + + is_invalid = indices == -1 + if HAS_TOPK_LENGTH: + is_invalid = is_invalid | (offs_n >= topk_len) + + valid = mask_n & ~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_64 + 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 + ) + + valid_2d = valid[:, None] + + # Use helper function for KV processing + acc_0, acc_1, acc_2, acc_3, acc_4, acc_5, acc_6, acc_7, m_i, l_i = ( + _process_kv_block_aggressive( + kv_block_base, + nope_rope_offset, + scale_base_offset, + valid, + valid_2d, + q_0, + q_1, + q_2, + q_3, + q_4, + q_5, + q_6, + q_7, + acc_0, + acc_1, + acc_2, + acc_3, + acc_4, + acc_5, + acc_6, + acc_7, + m_i, + l_i, + offs_tile, + sm_scale, + TILE_SIZE, + D_NOPE, + LOG2E, + BLOCK_H, + BLOCK_N, + ) + ) + + lse = m_i + tl.math.log2(tl.where(l_i == 0.0, 1.0, l_i)) / LOG2E + is_lonely_q = l_i == 0.0 + + output_scale = tl.where(l_i == 0.0, 0.0, 1.0 / l_i) + acc_0 = tl.where(is_lonely_q[:, None], 0.0, acc_0 * output_scale[:, None]) + acc_1 = tl.where(is_lonely_q[:, None], 0.0, acc_1 * output_scale[:, None]) + acc_2 = tl.where(is_lonely_q[:, None], 0.0, acc_2 * output_scale[:, None]) + acc_3 = tl.where(is_lonely_q[:, None], 0.0, acc_3 * output_scale[:, None]) + acc_4 = tl.where(is_lonely_q[:, None], 0.0, acc_4 * output_scale[:, None]) + acc_5 = tl.where(is_lonely_q[:, None], 0.0, acc_5 * output_scale[:, None]) + acc_6 = tl.where(is_lonely_q[:, None], 0.0, acc_6 * output_scale[:, None]) + acc_7 = tl.where(is_lonely_q[:, None], 0.0, acc_7 * output_scale[:, None]) + lse = tl.where(is_lonely_q, float("+inf"), lse) + + 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 + row_ptrs = po_base + offs_h[:, None] * stride_po_h + + # Store partial output as float32 for better precision in combine kernel + tl.store(row_ptrs + offs_tile[None, :] * stride_po_d, acc_0, mask=mask_h[:, None]) + tl.store( + row_ptrs + (TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_1, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (2 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_2, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (3 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_3, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (4 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_4, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (5 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_5, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (6 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_6, + mask=mask_h[:, None], + ) + tl.store( + row_ptrs + (7 * TILE_SIZE + offs_tile[None, :]) * stride_po_d, + acc_7, + mask=mask_h[:, None], + ) + + 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, lse, mask=mask_h) + + +@triton.jit +def _combine_splitk_kernel( + PartialOutput, + PartialLSE, + AttnSink, + Output, + LSE, + total_tokens, + total_tokens_bucket, + h_q, + d_v, + stride_po_s, + stride_po_t, + stride_po_h, + stride_po_d, + stride_plse_s, + stride_plse_t, + stride_plse_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 kernel (SPLIT_K=4).""" + LOG2E: tl.constexpr = 1.4426950408889634 + NEG_INF = float("-inf") + POS_INF = float("+inf") + INF_THRESHOLD = 1e30 + + 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 + offs_d = tl.arange(0, BLOCK_D) + + stride_plse_s_64 = tl.cast(stride_plse_s, tl.int64) + stride_plse_t_64 = tl.cast(stride_plse_t, tl.int64) + + lse_0 = tl.load( + PartialLSE + + 0 * stride_plse_s_64 + + pid_t_64 * stride_plse_t_64 + + offs_h * stride_plse_h, + mask=mask_h, + other=POS_INF, + ) + lse_1 = tl.load( + PartialLSE + + 1 * stride_plse_s_64 + + pid_t_64 * stride_plse_t_64 + + offs_h * stride_plse_h, + mask=mask_h, + other=POS_INF, + ) + lse_2 = tl.load( + PartialLSE + + 2 * stride_plse_s_64 + + pid_t_64 * stride_plse_t_64 + + offs_h * stride_plse_h, + mask=mask_h, + other=POS_INF, + ) + lse_3 = tl.load( + PartialLSE + + 3 * stride_plse_s_64 + + pid_t_64 * stride_plse_t_64 + + offs_h * stride_plse_h, + mask=mask_h, + other=POS_INF, + ) + + lse_0_valid = tl.abs(lse_0) < INF_THRESHOLD + lse_1_valid = tl.abs(lse_1) < INF_THRESHOLD + lse_2_valid = tl.abs(lse_2) < INF_THRESHOLD + lse_3_valid = tl.abs(lse_3) < INF_THRESHOLD + + lse_0_safe = tl.where(lse_0_valid, lse_0, NEG_INF) + lse_1_safe = tl.where(lse_1_valid, lse_1, NEG_INF) + lse_2_safe = tl.where(lse_2_valid, lse_2, NEG_INF) + lse_3_safe = tl.where(lse_3_valid, lse_3, NEG_INF) + + max_lse = tl.maximum( + tl.maximum(lse_0_safe, lse_1_safe), tl.maximum(lse_2_safe, lse_3_safe) + ) + + exp_0 = tl.where(lse_0_valid, tl.math.exp2((lse_0_safe - max_lse) * LOG2E), 0.0) + exp_1 = tl.where(lse_1_valid, tl.math.exp2((lse_1_safe - max_lse) * LOG2E), 0.0) + exp_2 = tl.where(lse_2_valid, tl.math.exp2((lse_2_safe - max_lse) * LOG2E), 0.0) + exp_3 = tl.where(lse_3_valid, tl.math.exp2((lse_3_safe - max_lse) * LOG2E), 0.0) + + sum_exp = exp_0 + exp_1 + exp_2 + exp_3 + all_invalid = sum_exp == 0.0 + sum_exp_safe = tl.where(all_invalid, 1.0, sum_exp) + + combined_lse = max_lse + tl.math.log2(sum_exp_safe) / LOG2E + combined_lse = tl.where(all_invalid, POS_INF, combined_lse) + + if HAS_ATTN_SINK: + attn_sink_vals = tl.load(AttnSink + offs_h, mask=mask_h, other=0.0) + is_lonely = combined_lse > INF_THRESHOLD + lse_safe_for_sink = tl.where(is_lonely, 0.0, combined_lse) + diff = attn_sink_vals - lse_safe_for_sink + diff_clamped = tl.minimum(tl.maximum(diff, -100.0), 100.0) + exp_diff = tl.math.exp2(diff_clamped * LOG2E) + exp_diff = tl.where(is_lonely, 0.0, exp_diff) + denominator = 1.0 + exp_diff + sink_scale = 1.0 / denominator + sink_scale = tl.where(is_lonely, 1.0, sink_scale) + + scale_0 = (exp_0 / sum_exp_safe) * sink_scale + scale_1 = (exp_1 / sum_exp_safe) * sink_scale + scale_2 = (exp_2 / sum_exp_safe) * sink_scale + scale_3 = (exp_3 / sum_exp_safe) * sink_scale + else: + scale_0 = exp_0 / sum_exp_safe + scale_1 = exp_1 / sum_exp_safe + scale_2 = exp_2 / sum_exp_safe + scale_3 = exp_3 / sum_exp_safe + + scale_0 = tl.where(all_invalid, 0.0, scale_0) + scale_1 = tl.where(all_invalid, 0.0, scale_1) + scale_2 = tl.where(all_invalid, 0.0, scale_2) + scale_3 = tl.where(all_invalid, 0.0, scale_3) + + stride_po_s_64 = tl.cast(stride_po_s, tl.int64) + stride_po_t_64 = tl.cast(stride_po_t, tl.int64) + + po_base_0 = ( + PartialOutput + + 0 * stride_po_s_64 + + pid_t_64 * stride_po_t_64 + + offs_h[:, None] * stride_po_h + ) + po_base_1 = ( + PartialOutput + + 1 * stride_po_s_64 + + pid_t_64 * stride_po_t_64 + + offs_h[:, None] * stride_po_h + ) + po_base_2 = ( + PartialOutput + + 2 * stride_po_s_64 + + pid_t_64 * stride_po_t_64 + + offs_h[:, None] * stride_po_h + ) + po_base_3 = ( + PartialOutput + + 3 * stride_po_s_64 + + pid_t_64 * stride_po_t_64 + + offs_h[:, None] * stride_po_h + ) + + stride_o_t_64 = tl.cast(stride_o_t, tl.int64) + o_base = Output + pid_t_64 * stride_o_t_64 + offs_h[:, None] * stride_o_h + + for d_idx in range(4): + d_offs = d_idx * BLOCK_D + offs_d[None, :] + po_0 = tl.load( + po_base_0 + d_offs * stride_po_d, mask=mask_h[:, None], other=0.0 + ) + po_1 = tl.load( + po_base_1 + d_offs * stride_po_d, mask=mask_h[:, None], other=0.0 + ) + po_2 = tl.load( + po_base_2 + d_offs * stride_po_d, mask=mask_h[:, None], other=0.0 + ) + po_3 = tl.load( + po_base_3 + d_offs * stride_po_d, mask=mask_h[:, None], other=0.0 + ) + combined = ( + scale_0[:, None] * po_0 + + scale_1[:, None] * po_1 + + scale_2[:, None] * po_2 + + scale_3[:, None] * po_3 + ) + tl.store( + o_base + d_offs * stride_o_d, combined.to(tl.bfloat16), mask=mask_h[:, None] + ) + + stride_lse_t_64 = tl.cast(stride_lse_t, tl.int64) + lse_ptrs = LSE + pid_t_64 * stride_lse_t_64 + offs_h * stride_lse_h + tl.store(lse_ptrs, combined_lse, mask=mask_h) + + +@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=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), + ], + key=["total_tokens_bucket", "h_q", "d_v"], +) +@triton.jit +def _combine_splitk_kernel_8_optimized( + PartialOutput, + PartialLSE, + AttnSink, + Output, + LSE, + total_tokens, + total_tokens_bucket, + h_q, + d_v, + stride_po_s, + stride_po_t, + stride_po_h, + stride_po_d, + stride_plse_s, + stride_plse_t, + stride_plse_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, +): + """Optimized combine kernel for split-K=8 with autotuning for BLOCK_H.""" + LOG2E: tl.constexpr = 1.4426950408889634 + NEG_INF = float("-inf") + POS_INF = float("+inf") + INF_THRESHOLD = 1e30 + + 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 + offs_d = tl.arange(0, BLOCK_D) + + stride_plse_s_64 = tl.cast(stride_plse_s, tl.int64) + stride_plse_t_64 = tl.cast(stride_plse_t, tl.int64) + + # Load all 8 LSE values + lse_base = PartialLSE + pid_t_64 * stride_plse_t_64 + offs_h * stride_plse_h + lse_0 = tl.load(lse_base + 0 * stride_plse_s_64, mask=mask_h, other=POS_INF) + lse_1 = tl.load(lse_base + 1 * stride_plse_s_64, mask=mask_h, other=POS_INF) + lse_2 = tl.load(lse_base + 2 * stride_plse_s_64, mask=mask_h, other=POS_INF) + lse_3 = tl.load(lse_base + 3 * stride_plse_s_64, mask=mask_h, other=POS_INF) + lse_4 = tl.load(lse_base + 4 * stride_plse_s_64, mask=mask_h, other=POS_INF) + lse_5 = tl.load(lse_base + 5 * stride_plse_s_64, mask=mask_h, other=POS_INF) + lse_6 = tl.load(lse_base + 6 * stride_plse_s_64, mask=mask_h, other=POS_INF) + lse_7 = tl.load(lse_base + 7 * stride_plse_s_64, mask=mask_h, other=POS_INF) + + lse_0_valid = tl.abs(lse_0) < INF_THRESHOLD + lse_1_valid = tl.abs(lse_1) < INF_THRESHOLD + lse_2_valid = tl.abs(lse_2) < INF_THRESHOLD + lse_3_valid = tl.abs(lse_3) < INF_THRESHOLD + lse_4_valid = tl.abs(lse_4) < INF_THRESHOLD + lse_5_valid = tl.abs(lse_5) < INF_THRESHOLD + lse_6_valid = tl.abs(lse_6) < INF_THRESHOLD + lse_7_valid = tl.abs(lse_7) < INF_THRESHOLD + + lse_0_safe = tl.where(lse_0_valid, lse_0, NEG_INF) + lse_1_safe = tl.where(lse_1_valid, lse_1, NEG_INF) + lse_2_safe = tl.where(lse_2_valid, lse_2, NEG_INF) + lse_3_safe = tl.where(lse_3_valid, lse_3, NEG_INF) + lse_4_safe = tl.where(lse_4_valid, lse_4, NEG_INF) + lse_5_safe = tl.where(lse_5_valid, lse_5, NEG_INF) + lse_6_safe = tl.where(lse_6_valid, lse_6, NEG_INF) + lse_7_safe = tl.where(lse_7_valid, lse_7, NEG_INF) + + max_lse = tl.maximum( + tl.maximum( + tl.maximum(lse_0_safe, lse_1_safe), tl.maximum(lse_2_safe, lse_3_safe) + ), + tl.maximum( + tl.maximum(lse_4_safe, lse_5_safe), tl.maximum(lse_6_safe, lse_7_safe) + ), + ) + + exp_0 = tl.where(lse_0_valid, tl.math.exp2((lse_0_safe - max_lse) * LOG2E), 0.0) + exp_1 = tl.where(lse_1_valid, tl.math.exp2((lse_1_safe - max_lse) * LOG2E), 0.0) + exp_2 = tl.where(lse_2_valid, tl.math.exp2((lse_2_safe - max_lse) * LOG2E), 0.0) + exp_3 = tl.where(lse_3_valid, tl.math.exp2((lse_3_safe - max_lse) * LOG2E), 0.0) + exp_4 = tl.where(lse_4_valid, tl.math.exp2((lse_4_safe - max_lse) * LOG2E), 0.0) + exp_5 = tl.where(lse_5_valid, tl.math.exp2((lse_5_safe - max_lse) * LOG2E), 0.0) + exp_6 = tl.where(lse_6_valid, tl.math.exp2((lse_6_safe - max_lse) * LOG2E), 0.0) + exp_7 = tl.where(lse_7_valid, tl.math.exp2((lse_7_safe - max_lse) * LOG2E), 0.0) + + sum_exp = exp_0 + exp_1 + exp_2 + exp_3 + exp_4 + exp_5 + exp_6 + exp_7 + all_invalid = sum_exp == 0.0 + sum_exp_safe = tl.where(all_invalid, 1.0, sum_exp) + + combined_lse = max_lse + tl.math.log2(sum_exp_safe) / LOG2E + combined_lse = tl.where(all_invalid, POS_INF, combined_lse) + + if HAS_ATTN_SINK: + attn_sink_vals = tl.load(AttnSink + offs_h, mask=mask_h, other=0.0) + is_lonely = combined_lse > INF_THRESHOLD + lse_safe_for_sink = tl.where(is_lonely, 0.0, combined_lse) + diff = attn_sink_vals - lse_safe_for_sink + diff_clamped = tl.minimum(tl.maximum(diff, -100.0), 100.0) + exp_diff = tl.math.exp2(diff_clamped * LOG2E) + exp_diff = tl.where(is_lonely, 0.0, exp_diff) + denominator = 1.0 + exp_diff + sink_scale = 1.0 / denominator + sink_scale = tl.where(is_lonely, 1.0, sink_scale) + + scale_0 = (exp_0 / sum_exp_safe) * sink_scale + scale_1 = (exp_1 / sum_exp_safe) * sink_scale + scale_2 = (exp_2 / sum_exp_safe) * sink_scale + scale_3 = (exp_3 / sum_exp_safe) * sink_scale + scale_4 = (exp_4 / sum_exp_safe) * sink_scale + scale_5 = (exp_5 / sum_exp_safe) * sink_scale + scale_6 = (exp_6 / sum_exp_safe) * sink_scale + scale_7 = (exp_7 / sum_exp_safe) * sink_scale + else: + scale_0 = exp_0 / sum_exp_safe + scale_1 = exp_1 / sum_exp_safe + scale_2 = exp_2 / sum_exp_safe + scale_3 = exp_3 / sum_exp_safe + scale_4 = exp_4 / sum_exp_safe + scale_5 = exp_5 / sum_exp_safe + scale_6 = exp_6 / sum_exp_safe + scale_7 = exp_7 / sum_exp_safe + + scale_0 = tl.where(all_invalid, 0.0, scale_0) + scale_1 = tl.where(all_invalid, 0.0, scale_1) + scale_2 = tl.where(all_invalid, 0.0, scale_2) + scale_3 = tl.where(all_invalid, 0.0, scale_3) + scale_4 = tl.where(all_invalid, 0.0, scale_4) + scale_5 = tl.where(all_invalid, 0.0, scale_5) + scale_6 = tl.where(all_invalid, 0.0, scale_6) + scale_7 = tl.where(all_invalid, 0.0, scale_7) + + 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_t_64 * stride_po_t_64 + offs_h[:, None] * stride_po_h + po_base_0 = po_base + 0 * stride_po_s_64 + po_base_1 = po_base + 1 * stride_po_s_64 + po_base_2 = po_base + 2 * stride_po_s_64 + po_base_3 = po_base + 3 * stride_po_s_64 + po_base_4 = po_base + 4 * stride_po_s_64 + po_base_5 = po_base + 5 * stride_po_s_64 + po_base_6 = po_base + 6 * stride_po_s_64 + po_base_7 = po_base + 7 * stride_po_s_64 + + stride_o_t_64 = tl.cast(stride_o_t, tl.int64) + o_base = Output + pid_t_64 * stride_o_t_64 + offs_h[:, None] * stride_o_h + + # Loop over D dimension with BLOCK_D chunks + num_d_iters: tl.constexpr = (512 + BLOCK_D - 1) // BLOCK_D + for d_idx in tl.static_range(num_d_iters): + d_offs = d_idx * BLOCK_D + offs_d[None, :] + mask_d = d_offs < d_v + mask_hd = mask_h[:, None] & mask_d + + po_0 = tl.load(po_base_0 + d_offs * stride_po_d, mask=mask_hd, other=0.0) + po_1 = tl.load(po_base_1 + d_offs * stride_po_d, mask=mask_hd, other=0.0) + po_2 = tl.load(po_base_2 + d_offs * stride_po_d, mask=mask_hd, other=0.0) + po_3 = tl.load(po_base_3 + d_offs * stride_po_d, mask=mask_hd, other=0.0) + po_4 = tl.load(po_base_4 + d_offs * stride_po_d, mask=mask_hd, other=0.0) + po_5 = tl.load(po_base_5 + d_offs * stride_po_d, mask=mask_hd, other=0.0) + po_6 = tl.load(po_base_6 + d_offs * stride_po_d, mask=mask_hd, other=0.0) + po_7 = tl.load(po_base_7 + d_offs * stride_po_d, mask=mask_hd, other=0.0) + + combined = ( + scale_0[:, None] * po_0 + + scale_1[:, None] * po_1 + + scale_2[:, None] * po_2 + + scale_3[:, None] * po_3 + + scale_4[:, None] * po_4 + + scale_5[:, None] * po_5 + + scale_6[:, None] * po_6 + + scale_7[:, None] * po_7 + ) + tl.store(o_base + d_offs * stride_o_d, combined.to(tl.bfloat16), mask=mask_hd) + + stride_lse_t_64 = tl.cast(stride_lse_t, tl.int64) + lse_ptrs = LSE + pid_t_64 * stride_lse_t_64 + offs_h * stride_lse_h + tl.store(lse_ptrs, combined_lse, mask=mask_h) + + +@triton.jit +def _combine_splitk_kernel_2( + PartialOutput, + PartialLSE, + AttnSink, + Output, + LSE, + total_tokens, + total_tokens_bucket, + h_q, + d_v, + stride_po_s, + stride_po_t, + stride_po_h, + stride_po_d, + stride_plse_s, + stride_plse_t, + stride_plse_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 kernel (SPLIT_K=2).""" + LOG2E: tl.constexpr = 1.4426950408889634 + NEG_INF = float("-inf") + POS_INF = float("+inf") + INF_THRESHOLD = 1e30 + + 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 + offs_d = tl.arange(0, BLOCK_D) + + stride_plse_s_64 = tl.cast(stride_plse_s, tl.int64) + stride_plse_t_64 = tl.cast(stride_plse_t, tl.int64) + + lse_0 = tl.load( + PartialLSE + + 0 * stride_plse_s_64 + + pid_t_64 * stride_plse_t_64 + + offs_h * stride_plse_h, + mask=mask_h, + other=POS_INF, + ) + lse_1 = tl.load( + PartialLSE + + 1 * stride_plse_s_64 + + pid_t_64 * stride_plse_t_64 + + offs_h * stride_plse_h, + mask=mask_h, + other=POS_INF, + ) + + lse_0_valid = tl.abs(lse_0) < INF_THRESHOLD + lse_1_valid = tl.abs(lse_1) < INF_THRESHOLD + + lse_0_safe = tl.where(lse_0_valid, lse_0, NEG_INF) + lse_1_safe = tl.where(lse_1_valid, lse_1, NEG_INF) + + max_lse = tl.maximum(lse_0_safe, lse_1_safe) + + exp_0 = tl.where(lse_0_valid, tl.math.exp2((lse_0_safe - max_lse) * LOG2E), 0.0) + exp_1 = tl.where(lse_1_valid, tl.math.exp2((lse_1_safe - max_lse) * LOG2E), 0.0) + + sum_exp = exp_0 + exp_1 + all_invalid = sum_exp == 0.0 + sum_exp_safe = tl.where(all_invalid, 1.0, sum_exp) + + combined_lse = max_lse + tl.math.log2(sum_exp_safe) / LOG2E + combined_lse = tl.where(all_invalid, POS_INF, combined_lse) + + if HAS_ATTN_SINK: + attn_sink_vals = tl.load(AttnSink + offs_h, mask=mask_h, other=0.0) + is_lonely = combined_lse > INF_THRESHOLD + lse_safe_for_sink = tl.where(is_lonely, 0.0, combined_lse) + diff = attn_sink_vals - lse_safe_for_sink + diff_clamped = tl.minimum(tl.maximum(diff, -100.0), 100.0) + exp_diff = tl.math.exp2(diff_clamped * LOG2E) + exp_diff = tl.where(is_lonely, 0.0, exp_diff) + denominator = 1.0 + exp_diff + sink_scale = 1.0 / denominator + sink_scale = tl.where(is_lonely, 1.0, sink_scale) + + scale_0 = (exp_0 / sum_exp_safe) * sink_scale + scale_1 = (exp_1 / sum_exp_safe) * sink_scale + else: + scale_0 = exp_0 / sum_exp_safe + scale_1 = exp_1 / sum_exp_safe + + scale_0 = tl.where(all_invalid, 0.0, scale_0) + scale_1 = tl.where(all_invalid, 0.0, scale_1) + + stride_po_s_64 = tl.cast(stride_po_s, tl.int64) + stride_po_t_64 = tl.cast(stride_po_t, tl.int64) + + po_base_0 = ( + PartialOutput + + 0 * stride_po_s_64 + + pid_t_64 * stride_po_t_64 + + offs_h[:, None] * stride_po_h + ) + po_base_1 = ( + PartialOutput + + 1 * stride_po_s_64 + + pid_t_64 * stride_po_t_64 + + offs_h[:, None] * stride_po_h + ) + + stride_o_t_64 = tl.cast(stride_o_t, tl.int64) + o_base = Output + pid_t_64 * stride_o_t_64 + offs_h[:, None] * stride_o_h + + for d_idx in range(4): + d_offs = d_idx * BLOCK_D + offs_d[None, :] + po_0 = tl.load( + po_base_0 + d_offs * stride_po_d, mask=mask_h[:, None], other=0.0 + ) + po_1 = tl.load( + po_base_1 + d_offs * stride_po_d, mask=mask_h[:, None], other=0.0 + ) + combined = scale_0[:, None] * po_0 + scale_1[:, None] * po_1 + tl.store( + o_base + d_offs * stride_o_d, combined.to(tl.bfloat16), mask=mask_h[:, None] + ) + + stride_lse_t_64 = tl.cast(stride_lse_t, tl.int64) + lse_ptrs = LSE + pid_t_64 * stride_lse_t_64 + offs_h * stride_lse_h + tl.store(lse_ptrs, combined_lse, mask=mask_h) + + +def _select_split_k(topk: int, h_q: int, total_tokens: int = 64) -> int: + """Select optimal split_k based on topk, h_q, and total_tokens. + + The split_k parameter controls how many parallel splits are used to process + 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 + """ + if topk >= 8192: + return 4 + else: + return 2 + + +# ============================================================================ +# 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. + """ + + _buffers = {} + _device = None + + @classmethod + def get_buffers( + cls, split_k: int, total_tokens: int, h_q: int, d_v: int, device: torch.device + ): + """Get or create intermediate buffers for the given configuration.""" + key = (split_k, total_tokens, h_q, d_v, device) + + if key not in cls._buffers or cls._device != device: + cls._device = device + 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 + ) + + cls._buffers[key] = { + "partial_output": partial_output, + "partial_lse": partial_lse, + "stride_po": partial_output.stride(), + "stride_plse": partial_lse.stride(), + } + + return cls._buffers[key] + + @classmethod + def clear(cls): + """Clear all cached buffers.""" + cls._buffers.clear() + cls._device = None + + +def fused_gather_attn_decode_dsv4_dual_scope_low_overhead( + q: torch.Tensor, + kv_cache_main: torch.Tensor, + indices_main: torch.Tensor, + block_size_main: int, + kv_cache_extra: torch.Tensor, + indices_extra: torch.Tensor, + block_size_extra: int, + sm_scale: float, + topk_length_main: Optional[torch.Tensor] = None, + topk_length_extra: Optional[torch.Tensor] = None, + attn_sink: Optional[torch.Tensor] = None, + s_q: int = 1, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Low-overhead version of fused_gather_attn_decode_dsv4_dual_scope. + + This version uses pre-allocated intermediate buffers and cached strides + 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] + topk_extra = indices_extra.shape[1] + total_topk = topk_main + topk_extra + d_v = DSV4_D_V + device = q.device + + # Prepare main KV cache + kv_uint8_main = kv_cache_main.view(torch.uint8) + num_blocks_main = kv_cache_main.shape[0] + stride_kv_block_main = kv_uint8_main.stride(0) + kv_flat_main = kv_uint8_main.reshape(num_blocks_main, -1) + + # Prepare extra KV cache + kv_uint8_extra = kv_cache_extra.view(torch.uint8) + num_blocks_extra = kv_cache_extra.shape[0] + stride_kv_block_extra = kv_uint8_extra.stride(0) + kv_flat_extra = kv_uint8_extra.reshape(num_blocks_extra, -1) + + if q.dtype != torch.bfloat16 or not q.is_contiguous(): + q = q.to(torch.bfloat16).contiguous() + + if not indices_main.is_contiguous(): + indices_main = indices_main.contiguous() + if not indices_extra.is_contiguous(): + 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 + + 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 + ): + # Fall back to non-splitk version + return fused_gather_attn_decode_dsv4_dual_scope( + q, + kv_cache_main, + indices_main, + block_size_main, + kv_cache_extra, + indices_extra, + block_size_extra, + sm_scale, + topk_length_main, + topk_length_extra, + attn_sink, + s_q, + ) + + # 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 + buffers = SplitKBufferPool.get_buffers(split_k, total_tokens, h_q, d_v, device) + partial_output = buffers["partial_output"] + partial_lse = buffers["partial_lse"] + 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). + 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) + + # Prepare dummy tensors for optional parameters + topk_length_main_tensor = ( + topk_length_main if topk_length_main is not None else lse[:1, 0] + ) + topk_length_extra_tensor = ( + topk_length_extra if topk_length_extra is not None else lse[:1, 0] + ) + attn_sink_tensor = attn_sink if attn_sink is not None else lse[0, :] + + # Pre-compute strides + stride_q = q.stride() + stride_o = output.stride() + stride_lse = lse.stride() + + # Check if buffer ops should be disabled + kv_cache_size_main = stride_kv_block_main * num_blocks_main + kv_cache_size_extra = stride_kv_block_extra * num_blocks_extra + disable_buffer_ops = ( + kv_cache_size_main > BUFFER_OPS_DISABLE_THRESHOLD + or kv_cache_size_extra > BUFFER_OPS_DISABLE_THRESHOLD + ) + + # Grid for splitk kernel + grid_splitk = lambda meta: ( + triton.cdiv(h_q, meta["BLOCK_H"]), + total_tokens, + split_k, + ) + + # Run splitk kernel + if disable_buffer_ops: + with triton.knobs.amd.scope(): + triton.knobs.amd.use_buffer_ops = False + _fused_gather_attn_dsv4_dual_scope_splitk_kernel[grid_splitk]( + q, + kv_flat_main, + indices_main, + topk_length_main_tensor, + kv_flat_extra, + indices_extra, + topk_length_extra_tensor, + partial_output, + partial_lse, + sm_scale, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + topk_main, + num_blocks_main, + block_size_main, + topk_extra, + num_blocks_extra, + block_size_extra, + s_q, + topk_per_split, + stride_q[0], + stride_q[1], + stride_q[2], + stride_kv_block_main, + stride_kv_block_extra, + indices_main.stride(0), + indices_main.stride(1), + indices_extra.stride(0), + indices_extra.stride(1), + stride_po[0], + stride_po[1], + stride_po[2], + stride_po[3], + stride_plse[0], + stride_plse[1], + stride_plse[2], + HAS_TOPK_LENGTH_MAIN=topk_length_main is not None, + HAS_TOPK_LENGTH_EXTRA=topk_length_extra is not None, + ) + else: + _fused_gather_attn_dsv4_dual_scope_splitk_kernel[grid_splitk]( + q, + kv_flat_main, + indices_main, + topk_length_main_tensor, + kv_flat_extra, + indices_extra, + topk_length_extra_tensor, + partial_output, + partial_lse, + sm_scale, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + topk_main, + num_blocks_main, + block_size_main, + topk_extra, + num_blocks_extra, + block_size_extra, + s_q, + topk_per_split, + stride_q[0], + stride_q[1], + stride_q[2], + stride_kv_block_main, + stride_kv_block_extra, + indices_main.stride(0), + indices_main.stride(1), + indices_extra.stride(0), + indices_extra.stride(1), + stride_po[0], + stride_po[1], + stride_po[2], + stride_po[3], + stride_plse[0], + stride_plse[1], + stride_plse[2], + HAS_TOPK_LENGTH_MAIN=topk_length_main is not None, + HAS_TOPK_LENGTH_EXTRA=topk_length_extra is not None, + ) + + # Run combine kernel + if split_k == 8: + grid_combine = lambda meta: (total_tokens, triton.cdiv(h_q, meta["BLOCK_H"])) + _combine_splitk_kernel_8_optimized[grid_combine]( + partial_output, + partial_lse, + attn_sink_tensor, + output, + lse, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + d_v, + stride_po[0], + stride_po[1], + stride_po[2], + stride_po[3], + stride_plse[0], + stride_plse[1], + stride_plse[2], + stride_o[0], + stride_o[1], + stride_o[2], + stride_lse[0], + stride_lse[1], + HAS_ATTN_SINK=attn_sink is not None, + ) + else: + BLOCK_H_COMBINE = 16 + BLOCK_D_COMBINE = 128 + grid_combine = (total_tokens, triton.cdiv(h_q, BLOCK_H_COMBINE)) + + if split_k == 2: + combine_kernel = _combine_splitk_kernel_2 + elif split_k == 4: + combine_kernel = _combine_splitk_kernel + else: + raise ValueError(f"Unsupported split_k: {split_k}") + + combine_kernel[grid_combine]( + partial_output, + partial_lse, + attn_sink_tensor, + output, + lse, + total_tokens, + _bucket_total_tokens(total_tokens), + h_q, + d_v, + stride_po[0], + stride_po[1], + stride_po[2], + stride_po[3], + stride_plse[0], + stride_plse[1], + stride_plse[2], + stride_o[0], + stride_o[1], + stride_o[2], + stride_lse[0], + stride_lse[1], + HAS_ATTN_SINK=attn_sink is not None, + BLOCK_H=BLOCK_H_COMBINE, + BLOCK_D=BLOCK_D_COMBINE, + num_warps=4, + num_stages=1, + ) + + return output, lse 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 new file mode 100644 index 000000000..02891b91b --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_optimized.py @@ -0,0 +1,289 @@ +""" +Optimized Triton MLA Decode Kernels for DeepSeek V4. + +This module provides optimized sparse attention decode with reduced Python overhead. + +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 + +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 ( + fused_gather_attn_decode_dsv4, + fused_gather_attn_decode_dsv4_dual_scope_low_overhead, +) + + +def triton_sparse_attn_decode( + 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]: + """Optimized sparse attention decode for DeepSeek V4 (d_qk=512).""" + d_qk = q.shape[-1] + + if d_qk != DSV4_D_QK: + raise ValueError( + f"Unsupported d_qk: {d_qk}. Expected {DSV4_D_QK} (DeepSeek V4)" + ) + + return _triton_sparse_attn_decode_dsv4( + q, kv_scope, extra_kv_scope, sm_scale, d_v, attn_sink + ) + + +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. + + The fused kernel avoids allocating a large intermediate gathered_kv + buffer and eliminates a separate gather kernel launch. However, for + h_q > 64 with medium-to-large batch sizes and larger topk, the + non-splitk fused kernel suffers from low GPU utilization (the grid + has only cdiv(h_q, BLOCK_H) blocks in the H dimension). In those + cases the fallback (separate gather + attention) can be faster on + the GPU, though it incurs extra torch.empty() overhead in CUDA + graphs. + + 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). + # For small topk (c128 layers, topk~192), fused always wins. + # For larger topk (c4 layers, topk~640), fused wins at small bs + # but the fallback catches up at bs>=16 due to better GPU utilization. + # However, the fallback has 4 extra torch.empty() calls that add + # ~30us CUDA-graph replay overhead, roughly cancelling the GPU gain. + # So we route to fused for all practical batch sizes. + if h_q > 64: + return total_tokens <= 256 + return True + + +def _triton_sparse_attn_decode_dsv4( + q: torch.Tensor, + kv_scope, + extra_kv_scope, + sm_scale: float, + d_v: int, + attn_sink: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Optimized sparse attention decode for DeepSeek V4 (d_qk=512).""" + 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 + block_size_main = kv_scope.blocked_k.shape[1] + + # 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 + + # Check if chunking needed (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) + 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() + + block_size_extra = extra_kv_scope.blocked_k.shape[1] + indices_extra = extra_kv_scope.indices_in_kvcache.reshape( + total_tokens, topk_extra + ) + if not indices_extra.is_contiguous(): + indices_extra = indices_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 new file mode 100644 index 000000000..2f6c6e789 --- /dev/null +++ b/python/sglang/srt/layers/attention/nsa/triton_decode/triton_mla_kernels_decode_splitk.py @@ -0,0 +1,534 @@ +""" +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 diff --git a/python/sglang/srt/layers/deepseek_v4_rope.py b/python/sglang/srt/layers/deepseek_v4_rope.py index 1b6c86698..69a8da7bf 100644 --- a/python/sglang/srt/layers/deepseek_v4_rope.py +++ b/python/sglang/srt/layers/deepseek_v4_rope.py @@ -288,6 +288,92 @@ def _fused_norm_rope_kernel( ) +@triton.jit +def _fused_softmax_pool_kernel( + kv_score_ptr, + out_ptr, + stride_bs: tl.constexpr, + stride_k: tl.constexpr, + K: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + base = pid * stride_bs + + offs = tl.arange(0, HEAD_BLOCK) + mask = offs < HEAD_DIM + + max_val = tl.full([HEAD_BLOCK], float("-inf"), dtype=tl.float32) + for k in range(K): + s = tl.load( + kv_score_ptr + base + k * stride_k + HEAD_DIM + offs, + mask=mask, + other=float("-inf"), + ).to(tl.float32) + max_val = tl.maximum(max_val, s) + + sum_exp = tl.zeros([HEAD_BLOCK], dtype=tl.float32) + weighted = tl.zeros([HEAD_BLOCK], dtype=tl.float32) + for k in range(K): + s = tl.load( + kv_score_ptr + base + k * stride_k + HEAD_DIM + offs, + mask=mask, + other=float("-inf"), + ).to(tl.float32) + v = tl.load( + kv_score_ptr + base + k * stride_k + offs, + mask=mask, + other=0.0, + ).to(tl.float32) + w = tl.exp(s - max_val) + sum_exp += w + weighted += v * w + + result = weighted / sum_exp + tl.store( + out_ptr + pid * HEAD_DIM + offs, result.to(out_ptr.dtype.element_ty), mask=mask + ) + + +def fused_softmax_pool_triton( + kv_score: torch.Tensor, + head_dim: int, +) -> torch.Tensor: + """Fused softmax-weighted-sum: out = (kv * softmax(score, dim=1)).sum(dim=1). + + Replaces the generic cunn_SpatialSoftMaxForward + elementwise multiply + sum + with a single Triton kernel. + + Args: + kv_score: [bs, K, 2 * head_dim] where first head_dim is kv, second is score. + head_dim: dimension of each of kv and score. + Returns: + output: [bs, head_dim] + """ + assert kv_score.dim() == 3 + bs, K, last = kv_score.shape + assert last == 2 * head_dim + assert kv_score.is_contiguous() + + out = torch.empty(bs, head_dim, dtype=kv_score.dtype, device=kv_score.device) + if bs == 0: + return out + + HEAD_BLOCK = triton.next_power_of_2(head_dim) + grid = (bs,) + _fused_softmax_pool_kernel[grid]( + kv_score, + out, + stride_bs=kv_score.stride(0), + stride_k=kv_score.stride(1), + K=K, + HEAD_DIM=head_dim, + HEAD_BLOCK=HEAD_BLOCK, + ) + return out + + def fused_norm_rope_inplace_triton( kv: torch.Tensor, weight: Optional[torch.Tensor], diff --git a/python/sglang/srt/layers/fused_qk_norm.py b/python/sglang/srt/layers/fused_qk_norm.py new file mode 100644 index 000000000..ce4bc0e42 --- /dev/null +++ b/python/sglang/srt/layers/fused_qk_norm.py @@ -0,0 +1,157 @@ +"""Fused Q/K RMSNorm in a single Triton kernel launch. + +Ported from ATOM (atom/model_ops/layernorm.py). Fuses per-head Q RMSNorm +(optionally weightless) and KV RMSNorm into one kernel, halving the number +of norm kernel launches per attention layer. +""" + +from typing import Optional, Tuple + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fused_qk_norm_kernel( + q_ptr, + k_ptr, + q_out_ptr, + k_out_ptr, + q_weight_ptr, + k_weight_ptr, + eps, + num_tokens, + head_dim, + q_in_stride0, + k_in_stride0, + q_out_stride0, + k_out_stride0, + num_q_heads, + num_k_heads, + Q_HAS_WEIGHT: tl.constexpr, + RBLOCK: tl.constexpr, + XBLOCK: tl.constexpr, +): + num_q_rows = num_tokens * num_q_heads + total_rows = num_tokens * (num_q_heads + num_k_heads) + + xoffset = tl.program_id(0) * XBLOCK + xindex = xoffset + tl.arange(0, XBLOCK)[:, None] + xmask = xindex < total_rows + cols = tl.arange(0, RBLOCK)[None, :] + col_mask = cols < head_dim + + is_q = xindex < num_q_rows + row_in_section = tl.where(is_q, xindex, xindex - num_q_rows) + cur_num_heads = tl.where(is_q, num_q_heads, num_k_heads) + + tokens = row_in_section // cur_num_heads + heads = row_in_section % cur_num_heads + + in_stride = tl.where(is_q, q_in_stride0, k_in_stride0) + in_bases = tokens * in_stride + heads * head_dim + + out_stride0 = tl.where(is_q, q_out_stride0, k_out_stride0) + out_bases = tokens * out_stride0 + heads * head_dim + + mask = xmask & col_mask + + if Q_HAS_WEIGHT: + qw = tl.load( + q_weight_ptr + cols, mask=col_mask, other=0.0, eviction_policy="evict_last" + ).to(tl.float32) + else: + qw = tl.full((RBLOCK,), 1.0, tl.float32) + kw = tl.load( + k_weight_ptr + cols, mask=col_mask, other=0.0, eviction_policy="evict_last" + ).to(tl.float32) + w = tl.where(is_q, qw, kw) + + x = tl.load( + q_ptr + in_bases + cols, + mask=mask & is_q, + other=0.0, + eviction_policy="evict_first", + ).to(tl.float32) + x = x + tl.load( + k_ptr + in_bases + cols, + mask=mask & ~is_q, + other=0.0, + eviction_policy="evict_first", + ).to(tl.float32) + + var = tl.sum(x * x, 1)[:, None] + rstd = tl.rsqrt(var / head_dim + eps) + + out = (x * rstd * w).to(q_out_ptr.dtype.element_ty) + tl.store( + q_out_ptr + out_bases + cols, + out, + mask=mask & is_q, + eviction_policy="evict_first", + ) + tl.store( + k_out_ptr + out_bases + cols, + out, + mask=mask & ~is_q, + eviction_policy="evict_first", + ) + + +def fused_qk_norm( + q: torch.Tensor, + k: torch.Tensor, + q_weight: Optional[torch.Tensor], + k_weight: torch.Tensor, + eps: float, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fused Q/K RMSNorm in a single Triton kernel launch. + + Args: + q: [num_tokens, num_heads, head_dim] + k: [num_tokens, num_kv_heads, head_dim] + q_weight: [head_dim] norm weight, or None for weightless Q norm + k_weight: [head_dim] norm weight (always required) + eps: epsilon for numerical stability + + Returns: + (q_normed, k_normed) same shapes as inputs + """ + head_dim = k_weight.shape[0] + if q_weight is not None: + assert q_weight.shape[0] == head_dim + num_tokens = q.shape[0] + num_q_heads = q.shape[1] + num_k_heads = k.shape[1] + total_rows = num_tokens * (num_q_heads + num_k_heads) + RBLOCK = triton.next_power_of_2(head_dim) + + q_out = torch.empty_like(q) + k_out = torch.empty_like(k) + + XBLOCK = 2 if total_rows > 8192 else 1 + NUM_WARPS = 1 + q_weight_arg = q_weight if q_weight is not None else k_weight + _fused_qk_norm_kernel[((total_rows + XBLOCK - 1) // XBLOCK,)]( + q, + k, + q_out, + k_out, + q_weight_arg, + k_weight, + eps, + num_tokens, + head_dim, + q.stride(0), + k.stride(0), + q_out.stride(0), + k_out.stride(0), + num_q_heads, + num_k_heads, + Q_HAS_WEIGHT=q_weight is not None, + RBLOCK=RBLOCK, + XBLOCK=XBLOCK, + num_warps=NUM_WARPS, + ) + return q_out, k_out diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py index 0e4ab204c..ed402a283 100644 --- a/python/sglang/srt/layers/moe/moe_runner/aiter.py +++ b/python/sglang/srt/layers/moe/moe_runner/aiter.py @@ -56,6 +56,7 @@ class AiterMoeQuantInfo(MoeQuantInfo): doweight_stage1: bool = False hidden_pad: int = 0 intermediate_pad: int = 0 + swiglu_limit: float = 0.0 @dataclass @@ -116,6 +117,7 @@ class AiterRunnerCore(MoeRunnerCore): return AiterRunnerOutput(hidden_states=runner_input.hidden_states) from aiter.fused_moe import fused_moe + from aiter.ops.flydsl.moe_common import GateMode a1_scale = ( runner_input.a1_scale @@ -128,6 +130,9 @@ class AiterRunnerCore(MoeRunnerCore): extra["num_local_tokens"] = runner_input.num_local_tokens if runner_input.output_dtype is not None: extra["dtype"] = runner_input.output_dtype + if quant_info.swiglu_limit > 0: + extra["gate_mode"] = GateMode.INTERLEAVE.value + extra["swiglu_limit"] = quant_info.swiglu_limit output = fused_moe( hidden_states=runner_input.hidden_states, diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 3c073f370..d9127143a 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -898,20 +898,46 @@ def biased_topk_jit_kernel_impl( ): assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" - from sglang.jit_kernel.moe_fused_gate import moe_fused_gate + if _use_aiter and scoring_func == "sqrtsoftplus" and num_fused_shared_experts == 0: + from aiter import topk_gating - topk_weights, topk_ids = moe_fused_gate( - gating_output, - correction_bias, - topk=topk, - scoring_func=scoring_func, - num_fused_shared_experts=num_fused_shared_experts, - renormalize=renormalize, - routed_scaling_factor=routed_scaling_factor, - apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, - ) - topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32) - return topk_weights, topk_ids + num_tokens = gating_output.shape[0] + topk_weights = torch.empty( + (num_tokens, topk), dtype=torch.float32, device=gating_output.device + ) + topk_ids = torch.empty( + (num_tokens, topk), dtype=torch.int32, device=gating_output.device + ) + + topk_gating( + topk_weights, + topk_ids, + gating_output, + correction_bias, + renormalize, + routed_scaling_factor, + score_func="sqrtsoftplus", + ) + + return topk_weights, topk_ids + + else: + from sglang.jit_kernel.moe_fused_gate import moe_fused_gate + + topk_weights, topk_ids = moe_fused_gate( + gating_output, + correction_bias, + topk=topk, + scoring_func=scoring_func, + num_fused_shared_experts=num_fused_shared_experts, + renormalize=renormalize, + routed_scaling_factor=routed_scaling_factor, + apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, + ) + topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to( + torch.int32 + ) + return topk_weights, topk_ids @torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 78d666646..3fe0b1884 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -125,8 +125,11 @@ def _require_fp4_dtype(): if _use_aiter or _use_hip_int4: - from aiter.ops.shuffle import shuffle_weight - from aiter.utility.fp4_utils import e8m0_shuffle + from aiter.ops.shuffle import ( + shuffle_scale_a16w4, + shuffle_weight, + shuffle_weight_a16w4, + ) if _use_aiter: from sglang.srt.layers.quantization.fp8_utils import ( @@ -1217,8 +1220,10 @@ class Fp8MoEMethod(FusedMoEMethodBase): for scale_name in ("w13_weight_scale_inv", "w2_weight_scale_inv"): scale = getattr(layer, scale_name) num_experts, num_rows, _ = scale.shape - scale.data = e8m0_shuffle(scale.view(num_experts * num_rows, -1)).view( - num_experts, num_rows, -1 + # a8w4: aiter flydsl scale layout + is_w13_scale = scale_name == "w13_weight_scale_inv" + scale.data = shuffle_scale_a16w4( + scale.view(num_experts * num_rows, -1), num_experts, is_w13_scale ) layer.w13_weight.data = layer.w13_weight.data.view(fp4_weight_dtype) @@ -1226,11 +1231,12 @@ class Fp8MoEMethod(FusedMoEMethodBase): is_shuffled = _is_shuffle_moe_mxfp4 if is_shuffled: - layer.w13_weight.data = shuffle_weight( - layer.w13_weight.contiguous(), (16, 16) + # a8w4: aiter flydsl weight layout + layer.w13_weight.data = shuffle_weight_a16w4( + layer.w13_weight.contiguous(), 16, True ) - layer.w2_weight.data = shuffle_weight( - layer.w2_weight.contiguous(), (16, 16) + layer.w2_weight.data = shuffle_weight_a16w4( + layer.w2_weight.contiguous(), 16, False ) layer.w13_weight.is_shuffled = is_shuffled layer.w2_weight.is_shuffled = is_shuffled @@ -2075,6 +2081,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): w13_scale=w13_scale, w2_scale=w2_scale, expert_mask=layer.dispatcher.expert_mask_gpu if _use_aiter else None, + swiglu_limit=self.moe_runner_config.swiglu_limit or 0.0, ) diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index e3c399cf5..b44388f37 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -470,8 +470,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): enable_memory_saver, ) + indexer_size = ( + self.c4_logical_size + if (not _is_hip or envs.SGLANG_OPT_USE_COMPRESSOR_V2.get()) + else c4_size + ) self.c4_indexer_kv_pool = DeepSeekV4IndexerPool( - self.c4_logical_size if not _is_hip else c4_size, + indexer_size, c4_page_size, dtype, indexer_head_dim, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index e2656c0e4..8b254348a 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -840,8 +840,11 @@ class DeepseekV2MoE(nn.Module): **topk_kwargs, ) final_hidden_states = self.experts(hidden_states, topk_output) - if not (_is_cuda or _is_musa) or isinstance( - self.experts.quant_method, KTEPWrapperMethod + if ( + not _is_cuda + and not _is_musa + and not _use_aiter + or isinstance(self.experts.quant_method, KTEPWrapperMethod) ): final_hidden_states *= self.routed_scaling_factor diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index da6845ca0..2a89d22f4 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -536,6 +536,118 @@ class MQALayer(nn.Module): return q + def _forward_prepare_multi_stream_hip( + self, + x: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + attn_backend, + q_out: Optional[torch.Tensor] = None, + x_quant=None, + ) -> torch.Tensor: + """ATOM-style ROCm path: overlap compressors, keep Q/KV on main stream.""" + assert self.alt_streams is not None + assert len(self.alt_streams) >= 1 + + current_stream = torch.cuda.current_stream() + stream_compressor = self.alt_streams[0] + stream_indexer_compressor = ( + self.alt_streams[1] if len(self.alt_streams) > 1 else None + ) + + if self.compressor is not None: + stream_compressor.wait_stream(current_stream) + with torch.cuda.stream(stream_compressor): + attn_backend.forward_core_compressor( + x, forward_batch, self.layer_id, self.compressor + ) + + if self.indexer is not None and stream_indexer_compressor is not None: + stream_indexer_compressor.wait_stream(current_stream) + with torch.cuda.stream(stream_indexer_compressor): + attn_backend.forward_indexer_compressor( + x=x, + forward_batch=forward_batch, + layer_id=self.indexer.layer_id, + compressor=self.indexer.compressor, + ) + + x_linear = x_quant if x_quant is not None else x + if self.fuse_wqa_wkv: + qkv_a, _ = self.wqkv_a(x_linear) + q_lora = qkv_a[..., : self.q_lora_rank] + else: + q_lora, _ = self.wq_a(x_linear) + qkv_a = None + + if self.use_fused_qk_norm_rope: + if _is_gfx95_supported: + q_for_wqb, q_lora = _fused_rmsnorm_fp8_quant( + q_lora, + self.q_norm.weight, + self.q_norm.variance_epsilon, + ) + q, _ = self.wq_b(q_for_wqb) + else: + q_lora = self.q_norm(q_lora) + q, _ = self.wq_b(q_lora) + + kv = ( + qkv_a[..., self.q_lora_rank :] + if qkv_a is not None + else self.wkv(x_linear)[0] + ) + + from sglang.srt.layers.fused_qk_norm_rope_store import ( + fused_qk_norm_rope_swa_store, + ) + + token_to_kv_pool = get_token_to_kv_pool() + swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa( + forward_batch.out_cache_loc + ) + swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id] + swa_page_size = token_to_kv_pool.swa_kv_pool.page_size + + q = fused_qk_norm_rope_swa_store( + q=q, + kv=kv, + q_norm_weight=None, + kv_norm_weight=self.kv_norm.weight, + q_rms_eps=self.eps, + kv_rms_eps=self.eps, + rope_head_dim=self.qk_rope_head_dim, + cos_cache=self.cos_cache, + sin_cache=self.sin_cache, + positions=positions, + swa_cache=swa_cache, + swa_loc=swa_loc, + swa_page_size=swa_page_size, + q_out=q_out, + dtype=x.dtype, + ) + else: + q_lora = self.q_norm(q_lora) + q = self._compute_q_b(q_lora, positions, q_out) + self._compute_kv_to_cache(x_linear, positions, forward_batch, qkv_a=qkv_a) + + del qkv_a + + if self.indexer is not None: + current_stream.wait_stream(stream_compressor) + if stream_indexer_compressor is not None: + current_stream.wait_stream(stream_indexer_compressor) + self.indexer( + x=x, + q_lora=q_lora, + forward_batch=forward_batch, + skip_compressor=True, + ) + elif self.compressor is not None: + current_stream.wait_stream(stream_compressor) + + return q + def _forward_prepare( self, x: torch.Tensor, @@ -695,14 +807,24 @@ class MQALayer(nn.Module): if enable_multi_stream: # Multi-stream path always fuses cache write into the K kernel, # so the bf16 KV intermediate is gone. - q = self._forward_prepare_multi_stream( - x, - positions, - forward_batch, - attn_backend, - q_out, - x_quant=x_quant, - ) + if _is_hip: + q = self._forward_prepare_multi_stream_hip( + x, + positions, + forward_batch, + attn_backend, + q_out, + x_quant=x_quant, + ) + else: + q = self._forward_prepare_multi_stream( + x, + positions, + forward_batch, + attn_backend, + q_out, + x_quant=x_quant, + ) kv = None else: q, kv = self._forward_prepare( @@ -792,12 +914,20 @@ class DeepseekV4DecoderLayer(nn.Module): alt_streams=alt_streams, compress_ratio_override=compress_ratio_override, ) + moe_alt_stream = ( + alt_streams[0] + if ( + alt_streams is not None + and (_is_cuda or envs.SGLANG_ROCM_USE_MULTI_STREAM.get()) + ) + else None + ) self.mlp = deepseek_v2.DeepseekV2MoE( config=config, quant_config=moe_quant_config_override or quant_config, prefix=add_prefix("mlp", prefix), layer_id=self.layer_id, - alt_stream=alt_streams[0] if alt_streams is not None else None, + alt_stream=moe_alt_stream, is_nextn=is_nextn, is_deepseek_v4=True, ) @@ -1147,7 +1277,19 @@ class DeepseekV4Model(nn.Module): else: self.embed_tokens = PPMissingLayer() self.rms_norm_eps = config.rms_norm_eps - self.alt_streams = [torch.cuda.Stream() for _ in range(5)] if _is_cuda else None + use_stream_pool = _is_cuda or ( + _is_hip + and ( + envs.SGLANG_ROCM_USE_MULTI_STREAM.get() + or envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get() + ) + ) + num_alt_streams = 5 if _is_cuda else 2 + self.alt_streams = ( + [torch.cuda.Stream() for _ in range(num_alt_streams)] + if use_stream_pool + else None + ) self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, lambda idx, prefix: DeepseekV4DecoderLayer( diff --git a/sgl-kernel/benchmark/bench_dsv4_norm_rope.py b/sgl-kernel/benchmark/bench_dsv4_norm_rope.py new file mode 100644 index 000000000..3e52c415d --- /dev/null +++ b/sgl-kernel/benchmark/bench_dsv4_norm_rope.py @@ -0,0 +1,75 @@ +"""Benchmark for DeepSeek-V4 fused norm + RoPE kernels.""" + +import itertools + +import sgl_kernel +import torch +import triton +import triton.testing + +try: + from sglang.utils import is_in_ci + + IS_CI = is_in_ci() +except ImportError: + IS_CI = False + +batch_sizes = [1] if IS_CI else [1, 4, 16, 64, 256] +num_heads_list = [8] if IS_CI else [8, 16, 64] +head_dims = [192] if IS_CI else [128, 192] + +configs = list(itertools.product(batch_sizes, num_heads_list, head_dims)) + + +def torch_rmsnorm_rope( + q: torch.Tensor, freqs_cis: torch.Tensor, positions: torch.Tensor, eps: float +) -> torch.Tensor: + """Naive PyTorch reference: RMSNorm + RoPE.""" + rms = torch.sqrt(q.float().pow(2).mean(dim=-1, keepdim=True) + eps) + q_normed = (q.float() / rms).to(q.dtype) + return q_normed + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "num_heads", "head_dim"], + x_vals=configs, + line_arg="provider", + line_vals=["sglang", "torch"], + line_names=["SGL Kernel", "PyTorch"], + styles=[("green", "-"), ("red", "--")], + ylabel="µs (median)", + plot_name="dsv4-q-norm-rope-performance", + args={}, + ) +) +def benchmark_q_norm_rope(batch_size, num_heads, head_dim, provider): + torch.manual_seed(42) + eps = 1e-6 + max_pos = 8192 + rope_dim = 64 + + q_input = torch.randn( + batch_size, num_heads, head_dim, dtype=torch.bfloat16, device="cuda" + ) + q_output = torch.empty_like(q_input) + freqs_cis = torch.randn(max_pos, rope_dim, dtype=torch.float32, device="cuda") + positions = torch.randint( + 0, max_pos, (batch_size,), dtype=torch.int32, device="cuda" + ) + + if provider == "sglang": + fn = lambda: sgl_kernel.dsv4_fused_q_norm_rope( + q_input, freqs_cis, positions, eps, q_output + ) + else: + fn = lambda: torch_rmsnorm_rope(q_input, freqs_cis, positions, eps) + + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph( + fn, quantiles=[0.5, 0.2, 0.8] + ) + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +if __name__ == "__main__": + benchmark_q_norm_rope.run(print_data=True) diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index b7c01a083..b50687abc 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -214,6 +214,21 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def("apply_shuffle_mul_sum(Tensor input, Tensor output, Tensor permutation, Tensor? factors) -> ()"); m.impl("apply_shuffle_mul_sum", torch::kCUDA, &apply_shuffle_mul_sum); + // DeepSeek-V4 fused norm + rope + m.def( + "dsv4_fused_q_norm_rope(Tensor q_input, Tensor! q_output, Tensor freqs_cis, Tensor positions, float eps) -> ()"); + m.impl("dsv4_fused_q_norm_rope", torch::kCUDA, &dsv4_fused_q_norm_rope); + + m.def( + "dsv4_fused_k_norm_rope_flashmla(Tensor kv, Tensor kv_weight, Tensor freqs_cis, Tensor positions, " + "Tensor out_loc, Tensor! kvcache, float eps, int page_size) -> ()"); + m.impl("dsv4_fused_k_norm_rope_flashmla", torch::kCUDA, &dsv4_fused_k_norm_rope_flashmla); + + m.def( + "dsv4_fused_q_indexer_rope_hadamard_quant(Tensor q_input, Tensor! q_fp8, Tensor weight, " + "Tensor! weights_out, float weight_scale, Tensor freqs_cis, Tensor positions) -> ()"); + m.impl("dsv4_fused_q_indexer_rope_hadamard_quant", torch::kCUDA, &dsv4_fused_q_indexer_rope_hadamard_quant); + m.def( "fused_qk_norm_rope(Tensor! qkv, int num_heads_q, " "int num_heads_k, int num_heads_v, int head_dim, float eps, " diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index eb41ff5a3..d8833e80e 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -368,6 +368,35 @@ void apply_shuffle_mul_sum( const torch::Tensor& permutation, const std::optional& factors); +/* + * From csrc/elementwise (DeepSeek-V4 norm + rope) + */ +void dsv4_fused_q_norm_rope( + const at::Tensor& q_input, + at::Tensor& q_output, + const at::Tensor& freqs_cis, + const at::Tensor& positions, + double eps); + +void dsv4_fused_k_norm_rope_flashmla( + const at::Tensor& kv, + const at::Tensor& kv_weight, + const at::Tensor& freqs_cis, + const at::Tensor& positions, + const at::Tensor& out_loc, + at::Tensor& kvcache, + double eps, + int64_t page_size); + +void dsv4_fused_q_indexer_rope_hadamard_quant( + const at::Tensor& q_input, + at::Tensor& q_fp8, + const at::Tensor& weight, + at::Tensor& weights_out, + double weight_scale, + const at::Tensor& freqs_cis, + const at::Tensor& positions); + void fused_qk_norm_rope( torch::Tensor& qkv, int64_t num_heads_q, diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index 1b97271f2..ce6e9d049 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -36,6 +36,9 @@ else: concat_mla_absorb_q, concat_mla_k, copy_to_gpu_no_ce, + dsv4_fused_k_norm_rope_flashmla, + dsv4_fused_q_indexer_rope_hadamard_quant, + dsv4_fused_q_norm_rope, fused_add_rmsnorm, gelu_and_mul, gelu_tanh_and_mul, @@ -125,6 +128,7 @@ else: if torch.version.hip is not None: from sgl_kernel.elementwise import gelu_quick + from sgl_kernel.top_k import deepseek_v4_topk_transform_512 if hasattr(torch.version, "musa") and torch.version.musa is not None: from sgl_kernel.musa import ( @@ -152,6 +156,9 @@ else: "cutlass_mla_get_workspace_size", "dsv3_fused_a_gemm", "dsv3_router_gemm", + "dsv4_fused_k_norm_rope_flashmla", + "dsv4_fused_q_indexer_rope_hadamard_quant", + "dsv4_fused_q_norm_rope", "es_fp8_blockwise_scaled_grouped_mm", "es_sm100_mxfp8_blockscaled_grouped_mm", "es_sm100_mxfp8_blockscaled_grouped_quant", @@ -205,6 +212,7 @@ else: if torch.version.hip is not None: _DEBUG_EXPORT_NAMES.append("gelu_quick") + _DEBUG_EXPORT_NAMES.append("deepseek_v4_topk_transform_512") for _name in _DEBUG_EXPORT_NAMES: if _name in globals(): diff --git a/sgl-kernel/tests/test_dsv4_norm_rope.py b/sgl-kernel/tests/test_dsv4_norm_rope.py new file mode 100644 index 000000000..4dbbbdc05 --- /dev/null +++ b/sgl-kernel/tests/test_dsv4_norm_rope.py @@ -0,0 +1,130 @@ +"""Tests for DeepSeek-V4 fused norm + RoPE kernels.""" + +import math + +import pytest +import sgl_kernel +import torch + + +def _ref_rmsnorm_self(x: torch.Tensor, eps: float) -> torch.Tensor: + """Reference: RMSNorm without weight (identity weight).""" + rms = torch.sqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + eps) + return (x.float() / rms).to(x.dtype) + + +def _ref_rope_interleaved( + x: torch.Tensor, freqs_cis: torch.Tensor, positions: torch.Tensor, rope_dim: int +) -> torch.Tensor: + """Reference: apply RoPE to the last `rope_dim` elements (interleaved re/im).""" + out = x.clone() + B = x.size(0) + head_dim = x.size(-1) + nope_dim = head_dim - rope_dim + + for b in range(B): + pos = positions[b].item() + freq = freqs_cis[pos] # (rope_dim,) interleaved [re0, im0, re1, im1, ...] + rope_part = out[b, ..., nope_dim:].float() + # Reshape to pairs + pairs = rope_part.reshape(*rope_part.shape[:-1], rope_dim // 2, 2) + x_real = pairs[..., 0] + x_imag = pairs[..., 1] + freq_pairs = freq.reshape(rope_dim // 2, 2) + f_real = freq_pairs[:, 0] + f_imag = freq_pairs[:, 1] + rot_real = x_real * f_real - x_imag * f_imag + rot_imag = x_real * f_imag + x_imag * f_real + result = torch.stack([rot_real, rot_imag], dim=-1).reshape(rope_part.shape) + out[b, ..., nope_dim:] = result.to(x.dtype) + return out + + +@pytest.mark.parametrize("batch_size", [1, 4, 16]) +@pytest.mark.parametrize("num_heads", [1, 8]) +@pytest.mark.parametrize("head_dim", [128, 192]) +def test_fused_q_norm_rope_correctness(batch_size, num_heads, head_dim): + """Test Q norm + rope against reference.""" + torch.manual_seed(42) + rope_dim = 64 + max_pos = 512 + eps = 1e-6 + + q_input = torch.randn( + batch_size, num_heads, head_dim, dtype=torch.bfloat16, device="cuda" + ) + freqs_cis = torch.randn(max_pos, rope_dim, dtype=torch.float32, device="cuda") + positions = torch.randint( + 0, max_pos, (batch_size,), dtype=torch.int32, device="cuda" + ) + + q_output = sgl_kernel.dsv4_fused_q_norm_rope(q_input, freqs_cis, positions, eps) + + # Reference + normed = _ref_rmsnorm_self(q_input, eps) + expected = _ref_rope_interleaved(normed, freqs_cis, positions, rope_dim) + + torch.testing.assert_close(q_output.float(), expected.float(), rtol=1e-2, atol=1e-2) + + +def test_fused_q_norm_rope_zero_batch(): + """Empty batch should not crash.""" + q_input = torch.empty(0, 8, 192, dtype=torch.bfloat16, device="cuda") + freqs_cis = torch.randn(512, 64, dtype=torch.float32, device="cuda") + positions = torch.empty(0, dtype=torch.int32, device="cuda") + q_output = sgl_kernel.dsv4_fused_q_norm_rope(q_input, freqs_cis, positions) + assert q_output.shape == q_input.shape + + +def test_fused_q_norm_rope_preallocated_output(): + """Test with pre-allocated output tensor.""" + torch.manual_seed(42) + B, H, D = 4, 8, 192 + q_input = torch.randn(B, H, D, dtype=torch.bfloat16, device="cuda") + freqs_cis = torch.randn(512, 64, dtype=torch.float32, device="cuda") + positions = torch.randint(0, 512, (B,), dtype=torch.int32, device="cuda") + q_output = torch.empty_like(q_input) + + result = sgl_kernel.dsv4_fused_q_norm_rope( + q_input, freqs_cis, positions, q_output=q_output + ) + assert result is q_output + + +@pytest.mark.parametrize("batch_size", [1, 8]) +def test_fused_q_indexer_rope_hadamard_quant_runs(batch_size): + """Smoke test: kernel runs without errors and produces finite results.""" + torch.manual_seed(42) + num_heads = 4 + head_dim = 128 + rope_dim = 64 + max_pos = 256 + + q_input = torch.randn( + batch_size, num_heads, head_dim, dtype=torch.bfloat16, device="cuda" + ) + q_fp8 = torch.empty( + batch_size, num_heads, head_dim, dtype=torch.uint8, device="cuda" + ) + weight = torch.randn(batch_size, num_heads, dtype=torch.bfloat16, device="cuda") + weights_out = torch.empty( + batch_size, num_heads, 1, dtype=torch.float32, device="cuda" + ) + freqs_cis = torch.randn(max_pos, rope_dim, dtype=torch.float32, device="cuda") + positions = torch.randint( + 0, max_pos, (batch_size,), dtype=torch.int32, device="cuda" + ) + weight_scale = 0.5 + + sgl_kernel.dsv4_fused_q_indexer_rope_hadamard_quant( + q_input, q_fp8, weight, weights_out, weight_scale, freqs_cis, positions + ) + + assert torch.isfinite(weights_out).all(), "weights_out contains non-finite values" + assert q_fp8.any(), "q_fp8 should not be all zeros" + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/manual/dsv4/test_fused_compress_attn_hip.py b/test/manual/dsv4/test_fused_compress_attn_hip.py new file mode 100644 index 000000000..0e6da091d --- /dev/null +++ b/test/manual/dsv4/test_fused_compress_attn_hip.py @@ -0,0 +1,465 @@ +"""Unit tests for the fused compressor attention Triton kernel on HIP. + +Validates numerical parity between the fused single-kernel path (plan-driven +Triton) and the reference per-seq Python implementation. + +Usage: + python -m pytest test/manual/dsv4/test_fused_compress_attn_hip.py -v + # or directly: + python test/manual/dsv4/test_fused_compress_attn_hip.py +""" + +import unittest +from dataclasses import dataclass + +import numpy as np +import torch + + +@dataclass +class FusedCompressPlan: + compress_plan_gpu: torch.Tensor + write_plan_gpu: torch.Tensor + num_compress: int + num_write: int + + +def write_current_token_to_state( + kv_score_input: torch.Tensor, + write_plan: torch.Tensor, + state_pool_buffer: torch.Tensor, + head_dim: int, + overlap: bool, + ratio: int, +) -> None: + """Reference write path used by this manual test. + + Plan row layout: [ragged_id, batch_id, position, window_len, state_base]. + """ + del head_dim # layout is already encoded in kv_score_input/state_pool_buffer shape. + state_size = (2 if overlap else 1) * ratio + plan_cpu = write_plan.cpu() + for row in plan_cpu: + ragged_id = int(row[0].item()) + position = int(row[2].item()) + state_base = int(row[4].item()) + if ragged_id < 0 or position < 0: + continue + dst = state_base + (position % state_size) + if ( + 0 <= dst < state_pool_buffer.shape[0] + and 0 <= ragged_id < kv_score_input.shape[0] + ): + state_pool_buffer[dst] = kv_score_input[ragged_id] + + +def fused_compress_attn( + state_pool_buffer: torch.Tensor, + plan: torch.Tensor, + ape: torch.Tensor, + rms_weight: torch.Tensor, + rms_eps: float, + freqs_cis_real: torch.Tensor, + head_dim: int, + rope_head_dim: int, + overlap: bool, + ratio: int, + out: torch.Tensor, +) -> torch.Tensor: + """Reference compress path for manual parity tests. + + This keeps the test runnable after removing `fused_compress_kernel.py`. + """ + freqs_cis = torch.view_as_complex( + freqs_cis_real.view(freqs_cis_real.shape[0], -1, 2).contiguous() + ) + result = _ref_compress( + kv_score_input=torch.empty( + 0, device=state_pool_buffer.device, dtype=torch.float32 + ), + state_pool=state_pool_buffer, + plan=plan, + ape=ape, + rms_weight=rms_weight, + rms_eps=rms_eps, + freqs_cis=freqs_cis, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + overlap=overlap, + ratio=ratio, + num_compress=plan.shape[0], + ) + out.copy_(result) + return out + + +def _make_plan_from_params( + extend_lens: list[int], + seq_lens: list[int], + ratio: int, + overlap: bool, + state_bases: list[int], + device: torch.device, +) -> FusedCompressPlan: + """Build a test plan without requiring real SWA / req_to_token tables.""" + bs = len(extend_lens) + ext = np.array(extend_lens, dtype=np.int32) + seq = np.array(seq_lens, dtype=np.int32) + total = int(ext.sum()) + + state_size = (2 if overlap else 1) * ratio + K = state_size + + batch_ids = np.repeat(np.arange(bs, dtype=np.int32), ext) + ragged_ids = np.arange(total, dtype=np.int32) + cu_extend = np.empty(bs + 1, dtype=np.int32) + cu_extend[0] = 0 + np.cumsum(ext, out=cu_extend[1:]) + j_in_seq = ragged_ids - cu_extend[batch_ids] + prefix_lens = seq - ext + positions = prefix_lens[batch_ids] + j_in_seq + + window_lens = np.maximum(0, K - np.minimum(j_in_seq + 1, K)).astype(np.int32) + state_base_arr = np.array(state_bases, dtype=np.int32) + state_base_per_token = state_base_arr[batch_ids] + + plan_rows = np.stack( + [ragged_ids, batch_ids, positions, window_lens, state_base_per_token], + axis=1, + ).astype(np.int32) + + compress_mask = (positions + 1) % ratio == 0 + compress_plan = plan_rows[compress_mask] + + write_starts = np.maximum(0, seq - K).astype(np.int32) + write_mask = positions >= write_starts[batch_ids] + write_plan = plan_rows[write_mask] + + n_compress = int(compress_plan.shape[0]) if compress_plan.size > 0 else 0 + n_write = int(write_plan.shape[0]) if write_plan.size > 0 else 0 + + compress_gpu = ( + torch.from_numpy(np.ascontiguousarray(compress_plan)).to(device) + if n_compress > 0 + else torch.empty((0, 5), dtype=torch.int32, device=device) + ) + write_gpu = ( + torch.from_numpy(np.ascontiguousarray(write_plan)).to(device) + if n_write > 0 + else torch.empty((0, 5), dtype=torch.int32, device=device) + ) + + return FusedCompressPlan( + compress_plan_gpu=compress_gpu, + write_plan_gpu=write_gpu, + num_compress=n_compress, + num_write=n_write, + ) + + +def _make_freqs_cis(max_seq: int, rope_dim: int, device: torch.device) -> torch.Tensor: + """Create test freqs_cis as complex64 [max_seq, rope_dim/2], matching production.""" + half = rope_dim // 2 + angles = torch.randn(max_seq, half, device=device, dtype=torch.float32) * 0.1 + return torch.polar(torch.ones_like(angles), angles) + + +def _freqs_to_real(freqs_cis: torch.Tensor) -> torch.Tensor: + """Convert complex64 freqs to float32 [max_seq, rope_dim] interleaved.""" + return torch.view_as_real(freqs_cis).flatten(-2).contiguous() + + +def _ref_compress( + kv_score_input: torch.Tensor, + state_pool: torch.Tensor, + plan: torch.Tensor, + ape: torch.Tensor, + rms_weight: torch.Tensor, + rms_eps: float, + freqs_cis: torch.Tensor, + head_dim: int, + rope_head_dim: int, + overlap: bool, + ratio: int, + num_compress: int, +) -> torch.Tensor: + """Pure-PyTorch reference matching SGLang compress_decode_paged semantics. + + State already has current tokens written (no APE). + APE is added to ALL K scores at compress time. + """ + if num_compress == 0: + return torch.empty(0, head_dim, dtype=torch.float32, device=state_pool.device) + + coff = 2 if overlap else 1 + half_dim = coff * head_dim + state_size = coff * ratio + K = state_size + + plan_cpu = plan[:num_compress].cpu() + out = torch.empty( + num_compress, head_dim, dtype=torch.float32, device=state_pool.device + ) + + for pid in range(num_compress): + position = int(plan_cpu[pid, 2].item()) + state_base = int(plan_cpu[pid, 4].item()) + + if position < 0: + continue + + kv_rows = [] + score_rows = [] + for k in range(K): + s = position - K + 1 + k + col_off = (head_dim if k >= ratio else 0) if overlap else 0 + ape_row = k % ratio + d_slice = slice(col_off, col_off + head_dim) + + if s < 0: + kv_rows.append( + torch.zeros(head_dim, dtype=torch.float32, device=state_pool.device) + ) + score_rows.append( + torch.full( + (head_dim,), + float("-inf"), + dtype=torch.float32, + device=state_pool.device, + ) + ) + else: + ring = s % state_size + row = state_pool[state_base + ring] + kv_rows.append(row[d_slice].float()) + # APE added to ALL scores + score_rows.append( + row[half_dim + col_off : half_dim + col_off + head_dim].float() + + ape[ape_row, d_slice].float() + ) + + kv_stack = torch.stack(kv_rows, dim=0) + sc_stack = torch.stack(score_rows, dim=0) + weights = torch.softmax(sc_stack, dim=0) + compressed = (weights * kv_stack).sum(dim=0) + + var = (compressed * compressed).mean() + normed = compressed * torch.rsqrt(var + rms_eps) * rms_weight.float() + + comp_pos = (position // ratio) * ratio + rope_seg = normed[-rope_head_dim:].clone() + freqs_row = torch.view_as_real(freqs_cis[comp_pos]).flatten() + cos_v = freqs_row[0::2].float() + sin_v = freqs_row[1::2].float() + + even = rope_seg[0::2] + odd = rope_seg[1::2] + normed[-rope_head_dim:] = torch.stack( + [even * cos_v - odd * sin_v, odd * cos_v + even * sin_v], dim=-1 + ).flatten() + + out[pid] = normed + + return out + + +class TestFusedCompressAttn(unittest.TestCase): + + def _run_test( + self, + ratio: int, + overlap: bool, + bs: int, + extend_lens: list[int], + prefix_lens: list[int], + head_dim: int = 512, + rope_head_dim: int = 64, + ): + device = torch.device("cuda") + torch.manual_seed(42) + coff = 2 if overlap else 1 + half_dim = coff * head_dim + last_dim = 2 * half_dim + state_size = coff * ratio + + seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)] + total_tokens = sum(extend_lens) + max_seq = max(seq_lens) + 128 + + kv_score_input = torch.randn( + total_tokens, last_dim, device=device, dtype=torch.float32 + ) + + pool_size = bs * state_size + 2 + state_pool = torch.randn( + pool_size, last_dim, device=device, dtype=torch.float32 + ) + state_pool[:, half_dim:] *= 0.5 # reasonable score magnitudes + + state_bases = [i * state_size for i in range(bs)] + ape = torch.randn(ratio, half_dim, device=device, dtype=torch.float32) * 0.1 + rms_weight = torch.ones(head_dim, device=device, dtype=torch.float32) + rms_eps = 1e-6 + freqs_cis = _make_freqs_cis(max_seq, rope_head_dim, device) + freqs_real = _freqs_to_real(freqs_cis) + + plan = _make_plan_from_params( + extend_lens, seq_lens, ratio, overlap, state_bases, device + ) + if plan.num_compress == 0: + return + + # Step 1: write current tokens to state (same for both paths) + state_triton = state_pool.clone() + state_ref = state_pool.clone() + + write_current_token_to_state( + kv_score_input=kv_score_input, + write_plan=plan.write_plan_gpu, + state_pool_buffer=state_triton, + head_dim=head_dim, + overlap=overlap, + ratio=ratio, + ) + # Reference: same write + write_current_token_to_state( + kv_score_input=kv_score_input, + write_plan=plan.write_plan_gpu, + state_pool_buffer=state_ref, + head_dim=head_dim, + overlap=overlap, + ratio=ratio, + ) + + # Step 2a: Triton fused compress + out_triton = torch.empty( + plan.num_compress, head_dim, device=device, dtype=torch.float32 + ) + fused_compress_attn( + state_pool_buffer=state_triton, + plan=plan.compress_plan_gpu, + ape=ape, + rms_weight=rms_weight, + rms_eps=rms_eps, + freqs_cis_real=freqs_real, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + overlap=overlap, + ratio=ratio, + out=out_triton, + ) + + # Step 2b: reference compress + out_ref = _ref_compress( + kv_score_input=kv_score_input, + state_pool=state_ref, + plan=plan.compress_plan_gpu, + ape=ape, + rms_weight=rms_weight, + rms_eps=rms_eps, + freqs_cis=freqs_cis, + head_dim=head_dim, + rope_head_dim=rope_head_dim, + overlap=overlap, + ratio=ratio, + num_compress=plan.num_compress, + ) + + torch.testing.assert_close(out_triton, out_ref, atol=1e-3, rtol=1e-3) + + def test_hca_single(self): + self._run_test( + ratio=128, overlap=False, bs=1, extend_lens=[128], prefix_lens=[0] + ) + + def test_hca_multi(self): + self._run_test( + ratio=128, overlap=False, bs=2, extend_lens=[128, 256], prefix_lens=[0, 128] + ) + + def test_csa_single(self): + self._run_test(ratio=4, overlap=True, bs=1, extend_lens=[16], prefix_lens=[8]) + + def test_csa_multi(self): + self._run_test( + ratio=4, overlap=True, bs=3, extend_lens=[8, 12, 16], prefix_lens=[4, 8, 0] + ) + + def test_csa_small_dim(self): + self._run_test( + ratio=4, + overlap=True, + bs=2, + extend_lens=[8, 8], + prefix_lens=[4, 0], + head_dim=256, + ) + + +class TestStateOrdering(unittest.TestCase): + + def test_write_then_compress(self): + """Verify write-first, compress-second matches reference.""" + device = torch.device("cuda") + torch.manual_seed(123) + ratio, overlap = 4, True + coff = 2 + head_dim, rope_head_dim = 128, 64 + half_dim = coff * head_dim + last_dim = 2 * half_dim + state_size = coff * ratio + + pool_size = state_size + 2 + state_pool = torch.randn( + pool_size, last_dim, device=device, dtype=torch.float32 + ) + state_pool[:, half_dim:] *= 0.5 + + kv_score_input = torch.randn(8, last_dim, device=device, dtype=torch.float32) + ape = torch.randn(ratio, half_dim, device=device, dtype=torch.float32) * 0.1 + rms_weight = torch.ones(head_dim, device=device, dtype=torch.float32) + freqs_cis = _make_freqs_cis(64, rope_head_dim, device) + + plan = _make_plan_from_params([8], [8], ratio, overlap, [0], device) + if plan.num_compress == 0: + return + + state_before = state_pool.clone() + + # Write first + write_current_token_to_state( + kv_score_input=kv_score_input, + write_plan=plan.write_plan_gpu, + state_pool_buffer=state_pool, + head_dim=head_dim, + overlap=overlap, + ratio=ratio, + ) + + # State should now be different (tokens written) + self.assertFalse(torch.allclose(state_pool, state_before)) + + # Compress + out = torch.empty( + plan.num_compress, head_dim, device=device, dtype=torch.float32 + ) + fused_compress_attn( + state_pool_buffer=state_pool, + plan=plan.compress_plan_gpu, + ape=ape, + rms_weight=rms_weight, + rms_eps=1e-6, + freqs_cis_real=_freqs_to_real(freqs_cis), + head_dim=head_dim, + rope_head_dim=rope_head_dim, + overlap=overlap, + ratio=ratio, + out=out, + ) + + self.assertFalse(torch.any(torch.isnan(out)).item()) + self.assertFalse(torch.any(torch.isinf(out)).item()) + + +if __name__ == "__main__": + unittest.main()