diff --git a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/__init__.py b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/__init__.py index f802d31f2..d7e34b6c9 100644 --- a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/__init__.py +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/__init__.py @@ -22,6 +22,14 @@ from __future__ import annotations +from tokenspeed_kernel_amd.ops.attention.gluon.dsa_gfx950 import ( # noqa: F401 + gluon_dsa_decode_gfx950, + gluon_dsa_prefill_gfx950, +) +from tokenspeed_kernel_amd.ops.attention.gluon.dsa_topk_gfx950 import ( # noqa: F401 + gluon_dsa_decode_topk_fp8_gfx950, + gluon_dsa_prefill_topk_fp8_gfx950, +) from tokenspeed_kernel_amd.ops.attention.gluon.mha_decode_gfx950 import ( # noqa: F401 gluon_mha_decode_gfx950, ) diff --git a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_gfx950.py b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_gfx950.py new file mode 100644 index 000000000..5996d9aed --- /dev/null +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_gfx950.py @@ -0,0 +1,1092 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""Selected-slot DSA Gluon kernels for AMD GFX950.""" + +from __future__ import annotations + +import torch +from tokenspeed_kernel_amd._triton import gl, gluon, tl, triton + +__all__ = [ + "gluon_dsa_decode_gfx950", + "gluon_dsa_prefill_gfx950", +] + +_REGISTERED_TOPK_WIDTHS = (512, 1024, 2048) + + +@gluon.constexpr_function +def _value_layout( + BLOCK_TOPK: gl.constexpr, + BLOCK_V: gl.constexpr, + NUM_WARPS: gl.constexpr, +): + return gl.BlockedLayout([1, 1], [1, 64], [NUM_WARPS, 1], [1, 0]) + + +@gluon.jit +def _dense_score( + q, + kv, + slots, + valid, + q_base, + kv_dim: gl.constexpr, + kv_lora_rank: gl.constexpr, + qk_rope_head_dim: gl.constexpr, + BLOCK_TOPK: gl.constexpr, + layout: gl.constexpr, +): + score = gl.full( + [BLOCK_TOPK], + value=0.0, + dtype=gl.float32, + layout=gl.SliceLayout(1, layout), + ) + for dim in gl.static_range(0, kv_lora_rank): + q_val = gl.load(q + q_base + dim).to(gl.float32) + k_val = gl.load( + kv + slots * kv_dim + dim, + mask=valid, + other=0.0, + ).to(gl.float32) + score += k_val * q_val + for dim in gl.static_range(0, qk_rope_head_dim): + q_val = gl.load(q + q_base + kv_lora_rank + dim).to(gl.float32) + k_val = gl.load( + kv + slots * kv_dim + kv_lora_rank + dim, + mask=valid, + other=0.0, + ).to(gl.float32) + score += k_val * q_val + return score + + +@gluon.jit +def _packed_score( + q, + kv_fp8, + kv_scale, + kv_rope, + slots, + valid, + q_base, + row_bytes: gl.constexpr, + kv_lora_rank: gl.constexpr, + qk_rope_head_dim: gl.constexpr, + BLOCK_TOPK: gl.constexpr, + layout: gl.constexpr, +): + score = gl.full( + [BLOCK_TOPK], + value=0.0, + dtype=gl.float32, + layout=gl.SliceLayout(1, layout), + ) + for dim in gl.static_range(0, kv_lora_rank): + q_val = gl.load(q + q_base + dim).to(gl.float32) + k_val = gl.load( + kv_fp8 + slots * row_bytes + dim, + mask=valid, + other=0.0, + ).to(gl.float32) + k_scale = gl.load( + kv_scale + (slots * row_bytes + kv_lora_rank + (dim // 128) * 4) // 4, + mask=valid, + other=0.0, + ).to(gl.float32) + score += k_val * k_scale * q_val + rope_base = (slots * row_bytes + kv_lora_rank + (kv_lora_rank // 128) * 4) // 2 + for dim in gl.static_range(0, qk_rope_head_dim): + q_val = gl.load(q + q_base + kv_lora_rank + dim).to(gl.float32) + k_val = gl.load( + kv_rope + rope_base + dim, + mask=valid, + other=0.0, + ).to(gl.float32) + score += k_val * q_val + return score + + +@gluon.jit +def _dsa_dense_mfma_kv_kernel( + q, + kv, + topk_indices, + topk_lens, + out, + stride_q_t: tl.int64, + stride_q_h: tl.int64, + stride_kv_t: tl.int64, + stride_o_t: tl.int64, + stride_o_h: tl.int64, + stride_topk_t: tl.int64, + scale: tl.float32, + num_heads: tl.int32, + TOPK: gl.constexpr, + BLOCK_H: gl.constexpr, + TILE_K: gl.constexpr, + D_V: gl.constexpr, + D_ROPE: gl.constexpr, +): + mfma_s: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + mfma_acc: gl.constexpr = gl.amd.cdna4.AMDMFMALayout( + version=4, + instr_shape=[16, 16, 16], + transposed=True, + warps_per_cta=[4, 1], + ) + + _qlora_tpw_k: gl.constexpr = min(64, D_V // 8) + _qlora_tpw_m: gl.constexpr = 64 // _qlora_tpw_k + blk_qlora: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[_qlora_tpw_m, _qlora_tpw_k], + warps_per_cta=[4, 1], + order=[1, 0], + ) + blk_qrope: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1, 8], + threads_per_warp=[8, 8], + warps_per_cta=[4, 1], + order=[1, 0], + ) + _klora_tpw_m: gl.constexpr = min(64, D_V // 8) + _klora_tpw_n: gl.constexpr = 64 // _klora_tpw_m + blk_klora: gl.constexpr = gl.BlockedLayout( + size_per_thread=[8, 1], + threads_per_warp=[_klora_tpw_m, _klora_tpw_n], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_krope: gl.constexpr = gl.BlockedLayout( + size_per_thread=[2, 1], + threads_per_warp=[32, 2], + warps_per_cta=[1, 4], + order=[0, 1], + ) + blk_topk: gl.constexpr = gl.BlockedLayout( + size_per_thread=[1], + threads_per_warp=[64], + warps_per_cta=[4], + order=[0], + ) + + sh_qlora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [BLOCK_H, D_V], + [1, 0], + ) + sh_qrope: gl.constexpr = gl.SwizzledSharedLayout( + vec=8, + per_phase=2, + max_phase=8, + order=[1, 0], + ) + sh_klora: gl.constexpr = gl.PaddedSharedLayout.with_identity_for( + [[512, 16]], + [D_V, TILE_K], + [0, 1], + ) + sh_krope: gl.constexpr = gl.SwizzledSharedLayout( + vec=8, + per_phase=2, + max_phase=8, + order=[0, 1], + ) + + dot_qlora_a: gl.constexpr = gl.DotOperandLayout( + operand_index=0, + parent=mfma_s, + k_width=8, + ) + dot_qrope_a: gl.constexpr = gl.DotOperandLayout( + operand_index=0, + parent=mfma_s, + k_width=8, + ) + dot_klora_b: gl.constexpr = gl.DotOperandLayout( + operand_index=1, + parent=mfma_s, + k_width=8, + ) + dot_krope_b: gl.constexpr = gl.DotOperandLayout( + operand_index=1, + parent=mfma_s, + k_width=8, + ) + dot_p_a: gl.constexpr = gl.DotOperandLayout( + operand_index=0, + parent=mfma_acc, + k_width=4, + ) + dot_v_b: gl.constexpr = gl.DotOperandLayout( + operand_index=1, + parent=mfma_acc, + k_width=4, + ) + + token_idx = gl.program_id(axis=0) + hg_idx = gl.program_id(axis=1) + hg_offset = hg_idx * BLOCK_H + valid_len = gl.load(topk_lens + token_idx).to(tl.int32) + + offs_h_qlora = hg_offset + gl.arange( + 0, + BLOCK_H, + layout=gl.SliceLayout(1, blk_qlora), + ) + offs_v_qlora = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_qlora = offs_h_qlora < num_heads + q_base = token_idx.to(tl.int64) * stride_q_t + q_offs_lora = ( + q_base + + offs_h_qlora[:, None].to(tl.int64) * stride_q_h + + offs_v_qlora[None, :].to(tl.int64) + ) + + offs_h_qrope = hg_offset + gl.arange( + 0, + BLOCK_H, + layout=gl.SliceLayout(1, blk_qrope), + ) + offs_r_qrope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(0, blk_qrope)) + mask_h_qrope = offs_h_qrope < num_heads + q_offs_rope = ( + q_base + + offs_h_qrope[:, None].to(tl.int64) * stride_q_h + + (D_V + offs_r_qrope[None, :]).to(tl.int64) + ) + + smem_qlora = gl.allocate_shared_memory( + q.dtype.element_ty, + [BLOCK_H, D_V], + layout=sh_qlora, + ) + smem_qrope = gl.allocate_shared_memory( + q.dtype.element_ty, + [BLOCK_H, D_ROPE], + layout=sh_qrope, + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qlora, + ptr=q, + offsets=q_offs_lora.to(tl.int32), + mask=mask_h_qlora[:, None], + ) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_qrope, + ptr=q, + offsets=q_offs_rope.to(tl.int32), + mask=mask_h_qrope[:, None], + ) + gl.amd.cdna4.async_copy.commit_group() + + NUM_TILES: gl.constexpr = (TOPK + TILE_K - 1) // TILE_K + topk_base = token_idx.to(tl.int64) * stride_topk_t + + offs_tile_klora = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_klora)) + offs_tile_krope = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, blk_krope)) + offs_tile_mma = gl.arange(0, TILE_K, layout=gl.SliceLayout(0, mfma_s)) + + offs_v_klora = gl.arange(0, D_V, layout=gl.SliceLayout(1, blk_klora)) + offs_r_krope = gl.arange(0, D_ROPE, layout=gl.SliceLayout(1, blk_krope)) + + topk_pos_klora = gl.amd.cdna4.buffer_load( + ptr=topk_indices, + offsets=topk_base.to(tl.int32) + offs_tile_klora, + mask=offs_tile_klora < TOPK, + other=-1, + ) + topk_pos_krope = gl.amd.cdna4.buffer_load( + ptr=topk_indices, + offsets=topk_base.to(tl.int32) + offs_tile_krope, + mask=offs_tile_krope < TOPK, + other=-1, + ) + topk_pos_mma = gl.amd.cdna4.buffer_load( + ptr=topk_indices, + offsets=topk_base.to(tl.int32) + offs_tile_mma, + mask=offs_tile_mma < TOPK, + other=-1, + ) + + valid_klora = (offs_tile_klora < valid_len) & (topk_pos_klora != -1) + valid_krope = (offs_tile_krope < valid_len) & (topk_pos_krope != -1) + valid_mma = (offs_tile_mma < valid_len) & (topk_pos_mma != -1) + safe_klora = gl.where(valid_klora, topk_pos_klora, 0) + safe_krope = gl.where(valid_krope, topk_pos_krope, 0) + + smem_krope = gl.allocate_shared_memory( + kv.dtype.element_ty, + [2, D_ROPE, TILE_K], + layout=sh_krope, + ) + smem_klora = gl.allocate_shared_memory( + kv.dtype.element_ty, + [2, D_V, TILE_K], + layout=sh_klora, + ) + + klora_offs = safe_klora[None, :].to(tl.int64) * stride_kv_t + offs_v_klora[ + :, None + ].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(0), + ptr=kv, + offsets=klora_offs.to(tl.int32), + mask=valid_klora[None, :], + ) + krope_offs = safe_krope[None, :].to(tl.int64) * stride_kv_t + ( + D_V + offs_r_krope[:, None] + ).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(0), + ptr=kv, + offsets=krope_offs.to(tl.int32), + mask=valid_krope[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + + gl.amd.cdna4.async_copy.wait_group(1) + q_lora_dot = smem_qlora.load(dot_qlora_a) + q_rope_dot = smem_qrope.load(dot_qrope_a) + + m_i = gl.full( + [BLOCK_H], + float("-inf"), + dtype=gl.float32, + layout=gl.SliceLayout(1, mfma_s), + ) + l_i = gl.full( + [BLOCK_H], + 0.0, + dtype=gl.float32, + layout=gl.SliceLayout(1, mfma_s), + ) + acc = gl.zeros([BLOCK_H, D_V], dtype=gl.float32, layout=mfma_acc) + + cur_buf = 0 + for tile_idx in range(NUM_TILES - 1): + next_base = (tile_idx + 1) * TILE_K + next_offs_klora = next_base + offs_tile_klora + next_offs_krope = next_base + offs_tile_krope + next_offs_mma = next_base + offs_tile_mma + + topk_pos_klora_next = gl.amd.cdna4.buffer_load( + ptr=topk_indices, + offsets=topk_base.to(tl.int32) + next_offs_klora, + mask=next_offs_klora < TOPK, + other=-1, + ) + topk_pos_krope_next = gl.amd.cdna4.buffer_load( + ptr=topk_indices, + offsets=topk_base.to(tl.int32) + next_offs_krope, + mask=next_offs_krope < TOPK, + other=-1, + ) + topk_pos_mma_next = gl.amd.cdna4.buffer_load( + ptr=topk_indices, + offsets=topk_base.to(tl.int32) + next_offs_mma, + mask=next_offs_mma < TOPK, + other=-1, + ) + + valid_klora_next = (next_offs_klora < valid_len) & (topk_pos_klora_next != -1) + valid_krope_next = (next_offs_krope < valid_len) & (topk_pos_krope_next != -1) + valid_mma_next = (next_offs_mma < valid_len) & (topk_pos_mma_next != -1) + safe_klora_next = gl.where(valid_klora_next, topk_pos_klora_next, 0) + safe_krope_next = gl.where(valid_krope_next, topk_pos_krope_next, 0) + next_buf = 1 - cur_buf + + klora_offs_next = safe_klora_next[None, :].to( + tl.int64 + ) * stride_kv_t + offs_v_klora[:, None].to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_klora.index(next_buf), + ptr=kv, + offsets=klora_offs_next.to(tl.int32), + mask=valid_klora_next[None, :], + ) + krope_offs_next = safe_krope_next[None, :].to(tl.int64) * stride_kv_t + ( + D_V + offs_r_krope[:, None] + ).to(tl.int64) + gl.amd.cdna4.async_copy.buffer_load_to_shared( + dest=smem_krope.index(next_buf), + ptr=kv, + offsets=krope_offs_next.to(tl.int32), + mask=valid_krope_next[None, :], + ) + gl.amd.cdna4.async_copy.commit_group() + gl.amd.cdna4.async_copy.wait_group(1) + + klora_smem_cur = smem_klora.index(cur_buf) + k_lora_t_dot = klora_smem_cur.load(dot_klora_b) + v_lora_dot = klora_smem_cur.permute([1, 0]).load(dot_v_b) + k_rope_t_dot = smem_krope.index(cur_buf).load(dot_krope_b) + + scores = gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + scores = gl.amd.cdna4.mfma(q_lora_dot, k_lora_t_dot, scores) + scores = gl.amd.cdna4.mfma(q_rope_dot, k_rope_t_dot, scores) + scores = scores * scale + + offs_h_mma = hg_offset + gl.arange( + 0, + BLOCK_H, + layout=gl.SliceLayout(1, mfma_s), + ) + mask_h_mma = offs_h_mma < num_heads + scores = gl.where( + valid_mma[None, :] & mask_h_mma[:, None], scores, -float("inf") + ) + + m_j = gl.max(scores, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > -float("inf"), m_new, 0.0) + alpha = gl.exp(m_i - m_new) + probs = gl.exp(scores - m_new[:, None]) + l_new = alpha * l_i + gl.sum(probs, axis=1) + + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + probs_dot = gl.convert_layout(probs.to(q.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(probs_dot, v_lora_dot, acc) + + m_i = m_new + l_i = l_new + cur_buf = next_buf + valid_mma = valid_mma_next + + gl.amd.cdna4.async_copy.wait_group(0) + + klora_smem_cur = smem_klora.index(cur_buf) + k_lora_t_dot = klora_smem_cur.load(dot_klora_b) + v_lora_dot = klora_smem_cur.permute([1, 0]).load(dot_v_b) + k_rope_t_dot = smem_krope.index(cur_buf).load(dot_krope_b) + + scores = gl.zeros([BLOCK_H, TILE_K], dtype=gl.float32, layout=mfma_s) + scores = gl.amd.cdna4.mfma(q_lora_dot, k_lora_t_dot, scores) + scores = gl.amd.cdna4.mfma(q_rope_dot, k_rope_t_dot, scores) + scores = scores * scale + + offs_h_mma = hg_offset + gl.arange( + 0, + BLOCK_H, + layout=gl.SliceLayout(1, mfma_s), + ) + mask_h_mma = offs_h_mma < num_heads + scores = gl.where(valid_mma[None, :] & mask_h_mma[:, None], scores, -float("inf")) + + m_j = gl.max(scores, axis=1) + m_new = gl.maximum(m_i, m_j) + m_new = gl.where(m_new > -float("inf"), m_new, 0.0) + alpha = gl.exp(m_i - m_new) + probs = gl.exp(scores - m_new[:, None]) + l_new = alpha * l_i + gl.sum(probs, axis=1) + + alpha_acc = gl.convert_layout(alpha, gl.SliceLayout(1, mfma_acc)) + acc = acc * alpha_acc[:, None] + probs_dot = gl.convert_layout(probs.to(q.dtype.element_ty), dot_p_a) + acc = gl.amd.cdna4.mfma(probs_dot, v_lora_dot, acc) + + l_i_acc = gl.convert_layout(l_new, gl.SliceLayout(1, mfma_acc)) + safe_l_i = gl.where(l_i_acc > 0.0, l_i_acc, 1.0) + acc = acc / safe_l_i[:, None] + acc = gl.where(l_i_acc[:, None] > 0.0, acc, 0.0) + + offs_h_o = hg_offset + gl.arange(0, BLOCK_H, layout=gl.SliceLayout(1, blk_qlora)) + offs_v_o = gl.arange(0, D_V, layout=gl.SliceLayout(0, blk_qlora)) + mask_h_o = offs_h_o < num_heads + o_base = token_idx.to(tl.int64) * stride_o_t + o_offs = ( + o_base + + offs_h_o[:, None].to(tl.int64) * stride_o_h + + offs_v_o[None, :].to(tl.int64) + ) + out_vals = gl.convert_layout(acc.to(out.dtype.element_ty), blk_qlora) + gl.amd.cdna4.buffer_store( + stored_value=out_vals, + ptr=out, + offsets=o_offs.to(tl.int32), + mask=mask_h_o[:, None], + ) + + +@gluon.jit +def _dsa_dense_kv_kernel( + q, + kv, + topk_indices, + topk_lens, + out, + num_heads: gl.constexpr, + head_dim: gl.constexpr, + kv_lora_rank: gl.constexpr, + qk_rope_head_dim: gl.constexpr, + kv_dim: gl.constexpr, + topk: gl.constexpr, + softmax_scale: gl.constexpr, + BLOCK_TOPK: gl.constexpr, + BLOCK_V: gl.constexpr, +): + token = gl.program_id(0) + head = gl.program_id(1) + v_block = gl.program_id(2) + layout: gl.constexpr = _value_layout(BLOCK_TOPK, BLOCK_V, gl.num_warps()) + topk_offsets = gl.arange(0, BLOCK_TOPK, layout=gl.SliceLayout(1, layout)) + v_offsets = v_block * BLOCK_V + gl.arange( + 0, BLOCK_V, layout=gl.SliceLayout(0, layout) + ) + q_base = (token * num_heads + head) * head_dim + valid_len = gl.load(topk_lens + token).to(gl.int32) + max_score = gl.full((), -float("inf"), gl.float32) + + for start in range(0, topk, BLOCK_TOPK): + cols = start + topk_offsets + valid = cols < valid_len + slots = gl.load(topk_indices + token * topk + cols, mask=valid, other=0).to( + gl.int64 + ) + valid = valid & (slots >= 0) + score = _dense_score( + q, + kv, + slots, + valid, + q_base, + kv_dim, + kv_lora_rank, + qk_rope_head_dim, + BLOCK_TOPK, + layout, + ) + score = gl.where(valid, score * softmax_scale, -float("inf")) + max_score = gl.maximum(max_score, gl.max(score, axis=0)) + + denom = gl.full((), 0.0, gl.float32) + acc = gl.full( + [BLOCK_V], + value=0.0, + dtype=gl.float32, + layout=gl.SliceLayout(0, layout), + ) + v_mask = v_offsets < kv_lora_rank + for start in range(0, topk, BLOCK_TOPK): + cols = start + topk_offsets + valid = cols < valid_len + slots = gl.load(topk_indices + token * topk + cols, mask=valid, other=0).to( + gl.int64 + ) + valid = valid & (slots >= 0) + score = _dense_score( + q, + kv, + slots, + valid, + q_base, + kv_dim, + kv_lora_rank, + qk_rope_head_dim, + BLOCK_TOPK, + layout, + ) + score = gl.where(valid, score * softmax_scale, -float("inf")) + probs = gl.exp(score - max_score) + probs = gl.where(valid, probs, 0.0) + denom += gl.sum(probs, axis=0) + v_vals = gl.load( + kv + slots[:, None] * kv_dim + v_offsets[None, :], + mask=valid[:, None] & v_mask[None, :], + other=0.0, + ).to(gl.float32) + acc += gl.sum(probs[:, None] * v_vals, axis=0) + + result = acc / denom + result = gl.where(denom > 0.0, result, 0.0) + out_base = (token * num_heads + head) * kv_lora_rank + gl.store(out + out_base + v_offsets, result, mask=v_mask) + + +@gluon.jit +def _dsa_packed_kv_kernel( + q, + kv_fp8, + kv_scale, + kv_rope, + topk_indices, + topk_lens, + out, + num_heads: gl.constexpr, + head_dim: gl.constexpr, + kv_lora_rank: gl.constexpr, + qk_rope_head_dim: gl.constexpr, + row_bytes: gl.constexpr, + topk: gl.constexpr, + softmax_scale: gl.constexpr, + BLOCK_TOPK: gl.constexpr, + BLOCK_V: gl.constexpr, +): + token = gl.program_id(0) + head = gl.program_id(1) + v_block = gl.program_id(2) + layout: gl.constexpr = _value_layout(BLOCK_TOPK, BLOCK_V, gl.num_warps()) + topk_offsets = gl.arange(0, BLOCK_TOPK, layout=gl.SliceLayout(1, layout)) + v_offsets = v_block * BLOCK_V + gl.arange( + 0, BLOCK_V, layout=gl.SliceLayout(0, layout) + ) + q_base = (token * num_heads + head) * head_dim + valid_len = gl.load(topk_lens + token).to(gl.int32) + max_score = gl.full((), -float("inf"), gl.float32) + + for start in range(0, topk, BLOCK_TOPK): + cols = start + topk_offsets + valid = cols < valid_len + slots = gl.load(topk_indices + token * topk + cols, mask=valid, other=0).to( + gl.int64 + ) + valid = valid & (slots >= 0) + score = _packed_score( + q, + kv_fp8, + kv_scale, + kv_rope, + slots, + valid, + q_base, + row_bytes, + kv_lora_rank, + qk_rope_head_dim, + BLOCK_TOPK, + layout, + ) + score = gl.where(valid, score * softmax_scale, -float("inf")) + max_score = gl.maximum(max_score, gl.max(score, axis=0)) + + denom = gl.full((), 0.0, gl.float32) + acc = gl.full( + [BLOCK_V], + value=0.0, + dtype=gl.float32, + layout=gl.SliceLayout(0, layout), + ) + v_mask = v_offsets < kv_lora_rank + for start in range(0, topk, BLOCK_TOPK): + cols = start + topk_offsets + valid = cols < valid_len + slots = gl.load(topk_indices + token * topk + cols, mask=valid, other=0).to( + gl.int64 + ) + valid = valid & (slots >= 0) + score = _packed_score( + q, + kv_fp8, + kv_scale, + kv_rope, + slots, + valid, + q_base, + row_bytes, + kv_lora_rank, + qk_rope_head_dim, + BLOCK_TOPK, + layout, + ) + score = gl.where(valid, score * softmax_scale, -float("inf")) + probs = gl.exp(score - max_score) + probs = gl.where(valid, probs, 0.0) + denom += gl.sum(probs, axis=0) + v_vals = gl.load( + kv_fp8 + slots[:, None] * row_bytes + v_offsets[None, :], + mask=valid[:, None] & v_mask[None, :], + other=0.0, + ).to(gl.float32) + v_scale = gl.load( + kv_scale + + ( + slots[:, None] * row_bytes + + kv_lora_rank + + (v_offsets[None, :] // 128) * 4 + ) + // 4, + mask=valid[:, None] & v_mask[None, :], + other=0.0, + ).to(gl.float32) + acc += gl.sum(probs[:, None] * v_vals * v_scale, axis=0) + + result = acc / denom + result = gl.where(denom > 0.0, result, 0.0) + out_base = (token * num_heads + head) * kv_lora_rank + gl.store(out + out_base + v_offsets, result, mask=v_mask) + + +def _flatten_packed_kv_cache(packed_kv_cache: torch.Tensor) -> torch.Tensor: + if packed_kv_cache.dim() == 2: + return packed_kv_cache + return packed_kv_cache.reshape(-1, packed_kv_cache.shape[-1]) + + +def _flatten_dense_kv_cache(kv_cache: torch.Tensor) -> torch.Tensor: + if kv_cache.dim() == 2: + return kv_cache + if kv_cache.dim() == 3: + return kv_cache.squeeze(1) + if kv_cache.shape[1] == 1: + kv_cache = kv_cache.permute(0, 2, 1, 3) + return kv_cache.reshape(-1, kv_cache.shape[-1]) + + +def _flatten_query(q: torch.Tensor) -> torch.Tensor: + if q.dim() == 3: + return q + return q.reshape(-1, q.shape[-2], q.shape[-1]) + + +def _check_inputs( + q: torch.Tensor, + topk_slots: torch.Tensor, + topk_lens: torch.Tensor | None, + *, + qk_nope_head_dim: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + page_size: int, +) -> None: + if q.dtype not in (torch.bfloat16, torch.float8_e4m3fn): + raise TypeError(f"Gluon DSA supports BF16/FP8 q, got {q.dtype}") + if page_size != 64: + raise ValueError(f"Gluon DSA supports page_size=64, got {page_size}") + if qk_nope_head_dim not in (128, 192): + raise ValueError( + "Gluon DSA supports qk_nope_head_dim in {128, 192}, got " + f"{qk_nope_head_dim}" + ) + if kv_lora_rank not in (128, 512): + raise ValueError( + f"Gluon DSA supports kv_lora_rank in {{128, 512}}, got {kv_lora_rank}" + ) + if qk_rope_head_dim != 64: + raise ValueError( + f"Gluon DSA supports qk_rope_head_dim=64, got {qk_rope_head_dim}" + ) + expected_head_dim = int(kv_lora_rank) + int(qk_rope_head_dim) + if q.shape[-1] != expected_head_dim: + raise ValueError( + "q head dim must equal kv_lora_rank + qk_rope_head_dim, got " + f"q={q.shape[-1]}, expected={expected_head_dim}" + ) + if topk_slots.dtype != torch.int32 or topk_slots.dim() != 2: + raise ValueError("topk_slots must be int32 with shape [tokens, topk]") + if topk_lens is None: + raise ValueError("Gluon DSA requires topk_lens for this milestone") + if topk_lens.dtype != torch.int32 or topk_lens.shape != (topk_slots.shape[0],): + raise ValueError("topk_lens must be int32 with shape [tokens]") + + +def _output_dtype(q: torch.Tensor) -> torch.dtype: + return torch.bfloat16 if q.dtype == torch.float8_e4m3fn else q.dtype + + +def _trim_topk_slots_for_context( + topk_slots: torch.Tensor, + max_seqlen_k: int, +) -> torch.Tensor: + topk = int(topk_slots.shape[1]) + effective_topk = topk + # Keep launch shapes on the widths registered by tokenspeed-kernel. + for supported_topk in _REGISTERED_TOPK_WIDTHS: + if max_seqlen_k <= supported_topk <= topk: + effective_topk = supported_topk + break + if effective_topk < int(topk_slots.shape[1]): + return topk_slots[:, :effective_topk].contiguous() + return topk_slots + + +def _run_dense_kv( + q: torch.Tensor, + kv_cache: torch.Tensor, + topk_slots: torch.Tensor, + topk_lens: torch.Tensor, + *, + softmax_scale: float, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> torch.Tensor: + out = torch.empty( + (q.shape[0], q.shape[1], kv_lora_rank), + dtype=_output_dtype(q), + device=q.device, + ) + grid = lambda meta: (q.shape[0], triton.cdiv(q.shape[1], meta["BLOCK_H"])) + _dsa_dense_mfma_kv_kernel[grid]( + q, + kv_cache, + topk_slots, + topk_lens, + out, + q.stride(0), + q.stride(1), + kv_cache.stride(0), + out.stride(0), + out.stride(1), + topk_slots.stride(0), + float(softmax_scale), + q.shape[1], + topk_slots.shape[1], + D_V=kv_lora_rank, + D_ROPE=qk_rope_head_dim, + BLOCK_H=16, + TILE_K=32, + num_warps=4, + ) + return out + + +def _run_dense_kv_scalar( + q: torch.Tensor, + kv_cache: torch.Tensor, + topk_slots: torch.Tensor, + topk_lens: torch.Tensor, + *, + softmax_scale: float, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> torch.Tensor: + kv_dim = int(kv_lora_rank) + int(qk_rope_head_dim) + out = torch.empty( + (q.shape[0], q.shape[1], kv_lora_rank), + dtype=_output_dtype(q), + device=q.device, + ) + _dsa_dense_kv_kernel[(q.shape[0], q.shape[1], triton.cdiv(kv_lora_rank, 64))]( + q, + kv_cache, + topk_slots, + topk_lens, + out, + q.shape[1], + q.shape[2], + kv_lora_rank, + qk_rope_head_dim, + kv_dim, + topk_slots.shape[1], + float(softmax_scale), + BLOCK_TOPK=32, + BLOCK_V=64, + num_warps=4, + ) + return out + + +def _run_packed_kv( + q: torch.Tensor, + packed_kv: torch.Tensor, + topk_slots: torch.Tensor, + topk_lens: torch.Tensor, + *, + softmax_scale: float, + kv_lora_rank: int, + qk_rope_head_dim: int, +) -> torch.Tensor: + row_bytes = int(packed_kv.shape[1]) + out = torch.empty( + (q.shape[0], q.shape[1], kv_lora_rank), + dtype=_output_dtype(q), + device=q.device, + ) + _dsa_packed_kv_kernel[(q.shape[0], q.shape[1], triton.cdiv(kv_lora_rank, 64))]( + q, + packed_kv.view(torch.float8_e4m3fn), + packed_kv.view(torch.float32), + packed_kv.view(torch.bfloat16), + topk_slots, + topk_lens, + out, + q.shape[1], + q.shape[2], + kv_lora_rank, + qk_rope_head_dim, + row_bytes, + topk_slots.shape[1], + float(softmax_scale), + BLOCK_TOPK=32, + BLOCK_V=64, + num_warps=4, + ) + return out + + +def _run_dsa( + *, + q: torch.Tensor, + kv_cache: torch.Tensor | None, + sparse_kv_cache: torch.Tensor | None, + topk_slots: torch.Tensor, + topk_lens: torch.Tensor, + qk_nope_head_dim: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + softmax_scale: float, + page_size: int, + k_scale: float, + out: torch.Tensor | None, + max_seqlen_k: int, +) -> torch.Tensor: + _check_inputs( + q, + topk_slots, + topk_lens, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + page_size=page_size, + ) + q = _flatten_query(q).contiguous() + topk_slots = topk_slots.contiguous() + topk_lens = topk_lens.contiguous() + topk_slots = _trim_topk_slots_for_context(topk_slots, max_seqlen_k) + softmax_scale = float(softmax_scale) * float(k_scale) + # The AITER-style tiled kernel maps to dense BF16 KV. TokenSpeed's packed + # sparse FP8 rows use a different physical layout and stay on the scalar path. + if sparse_kv_cache is not None: + result = _run_packed_kv( + q, + _flatten_packed_kv_cache(sparse_kv_cache).contiguous(), + topk_slots, + topk_lens, + softmax_scale=softmax_scale, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + ) + elif kv_cache is not None: + dense_kv = _flatten_dense_kv_cache(kv_cache).contiguous() + if ( + q.dtype == torch.bfloat16 + and dense_kv.dtype == torch.bfloat16 + and int(kv_lora_rank) == 512 + ): + result = _run_dense_kv( + q, + dense_kv, + topk_slots, + topk_lens, + softmax_scale=softmax_scale, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + ) + else: + result = _run_dense_kv_scalar( + q, + dense_kv, + topk_slots, + topk_lens, + softmax_scale=softmax_scale, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + ) + else: + raise ValueError("Gluon DSA requires kv_cache or sparse_kv_cache") + if out is None: + return result + out_view = out.reshape_as(result) + out_view.copy_(result) + return out + + +def gluon_dsa_decode_gfx950( + q: torch.Tensor, + kv_cache: torch.Tensor | None, + sparse_kv_cache: torch.Tensor | None, + topk_slots: torch.Tensor, + topk_lens: torch.Tensor | None, + max_seqlen_k: int, + qk_nope_head_dim: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + softmax_scale: float, + page_size: int, + q_len_per_req: int = 1, + logit_cap: float = 0.0, + k_scale: float = 1.0, + return_lse: bool = False, + out: torch.Tensor | None = None, +) -> torch.Tensor: + del q_len_per_req + if logit_cap != 0.0 or return_lse: + raise ValueError("Gluon DSA does not support logit_cap or return_lse") + return _run_dsa( + q=q, + kv_cache=kv_cache, + sparse_kv_cache=sparse_kv_cache, + topk_slots=topk_slots, + topk_lens=topk_lens, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + softmax_scale=softmax_scale, + page_size=page_size, + k_scale=k_scale, + out=out, + max_seqlen_k=max_seqlen_k, + ) + + +def gluon_dsa_prefill_gfx950( + q: torch.Tensor, + kv_cache: torch.Tensor | None, + sparse_kv_cache: torch.Tensor | None, + topk_slots: torch.Tensor, + topk_lens: torch.Tensor | None, + max_seqlen_k: int, + qk_nope_head_dim: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + softmax_scale: float, + page_size: int, + q_len_per_req: int = 1, + logit_cap: float = 0.0, + k_scale: float = 1.0, + return_lse: bool = False, + out: torch.Tensor | None = None, +) -> torch.Tensor: + del q_len_per_req + if logit_cap != 0.0 or return_lse: + raise ValueError("Gluon DSA does not support logit_cap or return_lse") + return _run_dsa( + q=q, + kv_cache=kv_cache, + sparse_kv_cache=sparse_kv_cache, + topk_slots=topk_slots, + topk_lens=topk_lens, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + softmax_scale=softmax_scale, + page_size=page_size, + k_scale=k_scale, + out=out, + max_seqlen_k=max_seqlen_k, + ) diff --git a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_score_gfx950.py b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_score_gfx950.py new file mode 100644 index 000000000..3a63da043 --- /dev/null +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_score_gfx950.py @@ -0,0 +1,266 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""DSA TopK scoring Gluon kernels for AMD GFX950.""" + +from __future__ import annotations + +import torch +from tokenspeed_kernel_amd._triton import gl, gluon + +__all__ = [ + "_check_packed_fp8_inputs", + "_dsa_decode_logits_fp8_kernel", + "_dsa_prefill_logits_fp8_kernel", +] + + +@gluon.constexpr_function +def _score_layout( + BLOCK_N: gl.constexpr, + BLOCK_D: gl.constexpr, + NUM_WARPS: gl.constexpr, +): + return gl.BlockedLayout([1, 8], [8, 8], [NUM_WARPS, 1], [1, 0]) + + +@gluon.jit +def _dsa_decode_logits_fp8_kernel( + q, + index_k_fp8, + index_k_scale, + weights, + seq_lens, + block_table, + logits, + block_table_stride: gl.constexpr, + logits_stride: gl.constexpr, + page_size: gl.constexpr, + row_bytes: gl.constexpr, + max_seq_len: gl.constexpr, + num_heads: gl.constexpr, + head_dim: gl.constexpr, + num_groups: gl.constexpr, + softmax_scale: gl.constexpr, + q_len_per_req: gl.constexpr, + BLOCK_N: gl.constexpr, + BLOCK_D: gl.constexpr, +): + token = gl.program_id(0) + block_id = gl.program_id(1) + layout: gl.constexpr = _score_layout(BLOCK_N, BLOCK_D, gl.num_warps()) + row_layout: gl.constexpr = gl.SliceLayout(1, layout) + dim_layout: gl.constexpr = gl.SliceLayout(0, layout) + offsets = block_id * BLOCK_N + gl.arange(0, BLOCK_N, layout=row_layout) + dim_offsets = gl.arange(0, BLOCK_D, layout=dim_layout) + req = token // q_len_per_req + q_offset = token - req * q_len_per_req + seq_len = gl.load(seq_lens + req).to(gl.int32) + if q_len_per_req != 1: + seq_len = seq_len - (q_len_per_req - 1) + q_offset + valid = (offsets < seq_len) & (offsets < max_seq_len) + block_idx = offsets // page_size + block_offset = offsets - block_idx * page_size + page = gl.amd.cdna4.buffer_load( + ptr=block_table, + offsets=(req * block_table_stride + block_idx).to(gl.int32), + mask=valid, + other=0, + ).to(gl.int64) + page_bytes = page_size * row_bytes + fp8_base = page * page_bytes + block_offset.to(gl.int64) * head_dim + scale_base = ( + page * (page_bytes // 4) + + (page_size * head_dim) // 4 + + block_offset.to(gl.int64) * num_groups + ) + scores = gl.full( + [BLOCK_N], + value=0.0, + dtype=gl.float32, + layout=row_layout, + ) + + for head in gl.static_range(0, num_heads): + head_weight = gl.load(weights + token * num_heads + head).to(gl.float32) + head_score = gl.full( + [BLOCK_N], + value=0.0, + dtype=gl.float32, + layout=row_layout, + ) + for dim_start in gl.static_range(0, head_dim, BLOCK_D): + dims = dim_start + dim_offsets + q_vals = gl.amd.cdna4.buffer_load( + ptr=q, + offsets=((token * num_heads + head) * head_dim + dims).to(gl.int32), + mask=dims < head_dim, + other=0.0, + ).to(gl.float32) + k_vals = gl.amd.cdna4.buffer_load( + ptr=index_k_fp8, + offsets=(fp8_base[:, None] + dims[None, :]).to(gl.int32), + mask=valid[:, None] & (dims[None, :] < head_dim), + other=0.0, + ).to(gl.float32) + k_scale = gl.amd.cdna4.buffer_load( + ptr=index_k_scale, + offsets=(scale_base + dim_start // 128).to(gl.int32), + mask=valid, + other=0.0, + ).to(gl.float32) + head_score += gl.sum(k_vals * k_scale[:, None] * q_vals[None, :], axis=1) + scores += head_score * head_weight + + scores *= softmax_scale + scores = gl.where(valid, scores, -float("inf")) + gl.store( + logits + token * logits_stride + offsets, + scores, + mask=offsets < max_seq_len, + ) + + +@gluon.jit +def _dsa_prefill_logits_fp8_kernel( + q, + index_k_fp8, + index_k_scale, + weights, + kv_workspace_slots, + row_starts, + row_ends, + logits, + logits_stride: gl.constexpr, + seq_len_sum: gl.constexpr, + page_size: gl.constexpr, + row_bytes: gl.constexpr, + num_heads: gl.constexpr, + head_dim: gl.constexpr, + num_groups: gl.constexpr, + softmax_scale: gl.constexpr, + BLOCK_N: gl.constexpr, + BLOCK_D: gl.constexpr, +): + token = gl.program_id(0) + block_id = gl.program_id(1) + layout: gl.constexpr = _score_layout(BLOCK_N, BLOCK_D, gl.num_warps()) + row_layout: gl.constexpr = gl.SliceLayout(1, layout) + dim_layout: gl.constexpr = gl.SliceLayout(0, layout) + offsets = block_id * BLOCK_N + gl.arange(0, BLOCK_N, layout=row_layout) + dim_offsets = gl.arange(0, BLOCK_D, layout=dim_layout) + row_start = gl.load(row_starts + token).to(gl.int32) + row_end = gl.load(row_ends + token).to(gl.int32) + valid = (offsets >= row_start) & (offsets < row_end) & (offsets < seq_len_sum) + slots = gl.amd.cdna4.buffer_load( + ptr=kv_workspace_slots, + offsets=offsets.to(gl.int32), + mask=offsets < seq_len_sum, + other=0, + ) + page = slots // page_size + block_offset = slots - page * page_size + page_bytes = page_size * row_bytes + fp8_base = page * page_bytes + block_offset * head_dim + scale_base = ( + page * (page_bytes // 4) + + (page_size * head_dim) // 4 + + block_offset * num_groups + ) + scores = gl.full( + [BLOCK_N], + value=0.0, + dtype=gl.float32, + layout=row_layout, + ) + + for head in gl.static_range(0, num_heads): + head_weight = gl.load(weights + token * num_heads + head).to(gl.float32) + head_score = gl.full( + [BLOCK_N], + value=0.0, + dtype=gl.float32, + layout=row_layout, + ) + for dim_start in gl.static_range(0, head_dim, BLOCK_D): + dims = dim_start + dim_offsets + q_vals = gl.amd.cdna4.buffer_load( + ptr=q, + offsets=((token * num_heads + head) * head_dim + dims).to(gl.int32), + mask=dims < head_dim, + other=0.0, + ).to(gl.float32) + k_vals = gl.amd.cdna4.buffer_load( + ptr=index_k_fp8, + offsets=(fp8_base[:, None] + dims[None, :]).to(gl.int32), + mask=valid[:, None] & (dims[None, :] < head_dim), + other=0.0, + ).to(gl.float32) + k_scale = gl.amd.cdna4.buffer_load( + ptr=index_k_scale, + offsets=(scale_base + dim_start // 128).to(gl.int32), + mask=valid, + other=0.0, + ).to(gl.float32) + head_score += gl.sum(k_vals * k_scale[:, None] * q_vals[None, :], axis=1) + scores += head_score * head_weight + + scores *= softmax_scale + scores = gl.where(valid, scores, -float("inf")) + gl.store( + logits + token * logits_stride + offsets, scores, mask=offsets < seq_len_sum + ) + + +def _check_packed_fp8_inputs( + q: torch.Tensor, + index_k_cache: torch.Tensor, + weights: torch.Tensor, + page_size: int, +) -> int: + if q.dtype != torch.bfloat16: + raise TypeError(f"DSA Gluon top-k expects BF16 q, got {q.dtype}") + if weights.dtype != torch.float32: + raise TypeError(f"DSA Gluon top-k expects FP32 weights, got {weights.dtype}") + if q.dim() != 3: + raise ValueError(f"q must be [tokens, heads, dim], got {tuple(q.shape)}") + if weights.shape != q.shape[:2]: + raise ValueError( + f"weights must have shape {tuple(q.shape[:2])}, got {tuple(weights.shape)}" + ) + if q.shape[2] != 128: + raise ValueError(f"DSA Gluon top-k supports head_dim=128, got {q.shape[2]}") + if page_size != 64: + raise ValueError(f"DSA Gluon top-k supports page_size=64, got {page_size}") + if index_k_cache.dtype != torch.uint8: + raise TypeError( + "DSA Gluon FP8 top-k expects uint8 packed index_k_cache, got " + f"{index_k_cache.dtype}" + ) + num_groups = q.shape[2] // 128 + row_bytes = q.shape[2] + num_groups * 4 + if index_k_cache.dim() != 2 or index_k_cache.shape[1] != row_bytes: + raise ValueError( + "packed index_k_cache must have shape [slots, row_bytes=" + f"{row_bytes}], got {tuple(index_k_cache.shape)}" + ) + if index_k_cache.shape[0] % page_size != 0: + raise ValueError("packed index_k_cache slot count must be page aligned") + return row_bytes diff --git a/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_topk_gfx950.py b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_topk_gfx950.py new file mode 100644 index 000000000..8ed370646 --- /dev/null +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_topk_gfx950.py @@ -0,0 +1,956 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + +"""DSA top-k Gluon kernels for AMD GFX950.""" + +from __future__ import annotations + +import torch +from tokenspeed_kernel_amd._triton import gl, gluon, triton +from tokenspeed_kernel_amd.ops.attention.gluon.dsa_score_gfx950 import ( + _check_packed_fp8_inputs, + _dsa_decode_logits_fp8_kernel, + _dsa_prefill_logits_fp8_kernel, +) + +_RADIX_TOPK_MIN_COLS = 65536 +_RADIX_TOPK_BLOCK_N = 4096 + +__all__ = [ + "gluon_dsa_decode_topk_fp8_gfx950", + "gluon_dsa_prefill_topk_fp8_gfx950", +] + + +@gluon.constexpr_function +def _vector_layout( + BLOCK: gl.constexpr, + NUM_WARPS: gl.constexpr, + LOAD_ELEMS: gl.constexpr, +): + return gl.BlockedLayout([LOAD_ELEMS], [64], [NUM_WARPS], [0]) + + +@gluon.jit +def _fp32_to_ordered_key(x): + bits = x.to(gl.uint32, bitcast=True) + sign = bits & 0x80000000 + return bits ^ gl.where(sign != 0, 0xFFFFFFFF, 0x80000000) + + +@gluon.jit +def _topk_add(a, b): + return a + b + + +@gluon.jit +def _find_topk_threshold_key( + values, + valid, + topk: gl.constexpr, + BLOCK_N: gl.constexpr, + layout: gl.constexpr, +): + keys = _fp32_to_ordered_key(values) + prefix = 0 + remaining = topk + + # Ordered FP32 keys are searched from the most-significant 4-bit nibble down. + for shift in gl.static_range(28, -1, -4): + if shift == 28: + prefix_match = valid + else: + prefix_match = valid & ((keys >> (shift + 4)) == prefix) + bucket = (keys >> shift) & 0xF + cumulative = 0 + selected = 0 + selected_remaining = remaining + found = 0 + + for bucket_id in gl.static_range(15, -1, -1): + in_bucket = prefix_match & (bucket == bucket_id) + count = gl.sum( + gl.where( + in_bucket, + gl.full([BLOCK_N], 1, gl.int32, layout=layout), + gl.full([BLOCK_N], 0, gl.int32, layout=layout), + ), + axis=0, + ).to(gl.int32) + take = (found == 0) & (remaining <= cumulative + count) + selected = gl.where(take, bucket_id, selected) + selected_remaining = gl.where( + take, remaining - cumulative, selected_remaining + ) + cumulative += gl.where(found == 0, count, 0) + found = gl.where(take, 1, found) + + prefix = (prefix << 4) | selected + remaining = selected_remaining + + return prefix + + +@gluon.jit +def _dsa_decode_select_topk_kernel( + logits, + block_table, + seq_lens, + out, + lens_out, + logits_stride: gl.constexpr, + block_table_stride: gl.constexpr, + out_stride: gl.constexpr, + block_table_cols: gl.constexpr, + page_size: gl.constexpr, + topk: gl.constexpr, + q_len_per_req: gl.constexpr, + BLOCK_N: gl.constexpr, + LOAD_ELEMS: gl.constexpr, + TOPK_LOAD_ELEMS: gl.constexpr, +): + row = gl.program_id(0) + layout: gl.constexpr = _vector_layout(BLOCK_N, gl.num_warps(), LOAD_ELEMS) + topk_layout: gl.constexpr = _vector_layout(topk, gl.num_warps(), TOPK_LOAD_ELEMS) + offsets = gl.arange(0, BLOCK_N, layout=layout) + top_offsets = gl.arange(0, topk, layout=topk_layout) + req = row // q_len_per_req + q_offset = row - req * q_len_per_req + seq_len = gl.load(seq_lens + req).to(gl.int32) + if q_len_per_req != 1: + seq_len = seq_len - (q_len_per_req - 1) + q_offset + lens = gl.minimum(seq_len, topk).to(gl.int32) + gl.store(lens_out + row, lens) + gl.store(out + row * out_stride + top_offsets, -1) + + if seq_len <= topk: + valid_top = top_offsets < seq_len + local = top_offsets.to(gl.int32) + block_idx = local // page_size + block_offset = local - block_idx * page_size + page = gl.load( + block_table + req * block_table_stride + block_idx, + mask=valid_top & (block_idx < block_table_cols), + other=0, + ).to(gl.int32) + slots = page * page_size + block_offset + gl.store( + out + row * out_stride + top_offsets, + gl.where(valid_top, slots, -1), + mask=top_offsets < topk, + ) + return + + valid = offsets < seq_len + values = gl.load( + logits + row * logits_stride + offsets, + mask=valid, + other=-float("inf"), + ) + threshold = _find_topk_threshold_key(values, valid, topk, BLOCK_N, layout) + keys = _fp32_to_ordered_key(values) + greater = valid & (keys > threshold) + equal = valid & (keys == threshold) + greater_i32 = greater.to(gl.int32) + equal_i32 = equal.to(gl.int32) + count_greater = gl.sum(greater_i32, axis=0).to(gl.int32) + greater_pos = gl.associative_scan(greater_i32, 0, _topk_add) - 1 + equal_pos = count_greater + gl.associative_scan(equal_i32, 0, _topk_add) - 1 + greater_write = greater & (greater_pos < topk) + equal_write = equal & (equal_pos < topk) + local = offsets.to(gl.int32) + block_idx = local // page_size + block_offset = local - block_idx * page_size + page = gl.load( + block_table + req * block_table_stride + block_idx, + mask=(greater_write | equal_write) & (block_idx < block_table_cols), + other=0, + ).to(gl.int32) + slots = page * page_size + block_offset + gl.store(out + row * out_stride + greater_pos, slots, mask=greater_write) + gl.store(out + row * out_stride + equal_pos, slots, mask=equal_write) + + +@gluon.jit +def _dsa_prefill_select_topk_kernel( + logits, + row_starts, + row_ends, + out, + lens_out, + logits_stride: gl.constexpr, + out_stride: gl.constexpr, + topk: gl.constexpr, + BLOCK_N: gl.constexpr, + LOAD_ELEMS: gl.constexpr, + TOPK_LOAD_ELEMS: gl.constexpr, +): + row = gl.program_id(0) + layout: gl.constexpr = _vector_layout(BLOCK_N, gl.num_warps(), LOAD_ELEMS) + topk_layout: gl.constexpr = _vector_layout(topk, gl.num_warps(), TOPK_LOAD_ELEMS) + offsets = gl.arange(0, BLOCK_N, layout=layout) + top_offsets = gl.arange(0, topk, layout=topk_layout) + row_start = gl.load(row_starts + row).to(gl.int32) + row_end = gl.load(row_ends + row).to(gl.int32) + candidate_len = gl.maximum(row_end - row_start, 0) + lens = gl.minimum(candidate_len, topk).to(gl.int32) + gl.store(lens_out + row, lens) + gl.store(out + row * out_stride + top_offsets, -1) + + if candidate_len <= topk: + local = row_start + top_offsets.to(gl.int32) + valid_top = top_offsets < candidate_len + gl.store( + out + row * out_stride + top_offsets, + gl.where(valid_top, local, -1), + mask=top_offsets < topk, + ) + return + + valid = (offsets >= row_start) & (offsets < row_end) + values = gl.load( + logits + row * logits_stride + offsets, + mask=valid, + other=-float("inf"), + ) + threshold = _find_topk_threshold_key(values, valid, topk, BLOCK_N, layout) + keys = _fp32_to_ordered_key(values) + greater = valid & (keys > threshold) + equal = valid & (keys == threshold) + greater_i32 = greater.to(gl.int32) + equal_i32 = equal.to(gl.int32) + count_greater = gl.sum(greater_i32, axis=0).to(gl.int32) + greater_pos = gl.associative_scan(greater_i32, 0, _topk_add) - 1 + equal_pos = count_greater + gl.associative_scan(equal_i32, 0, _topk_add) - 1 + greater_write = greater & (greater_pos < topk) + equal_write = equal & (equal_pos < topk) + local = offsets.to(gl.int32) + gl.store(out + row * out_stride + greater_pos, local, mask=greater_write) + gl.store(out + row * out_stride + equal_pos, local, mask=equal_write) + + +@gluon.jit +def _dsa_decode_radix_init_kernel( + seq_lens, + out, + lens_out, + prefixes, + remaining, + counters, + out_stride: gl.constexpr, + topk: gl.constexpr, + q_len_per_req: gl.constexpr, + TOPK_LOAD_ELEMS: gl.constexpr, +): + row = gl.program_id(0) + top_layout: gl.constexpr = _vector_layout(topk, gl.num_warps(), TOPK_LOAD_ELEMS) + top_offsets = gl.arange(0, topk, layout=top_layout) + req = row // q_len_per_req + q_offset = row - req * q_len_per_req + seq_len = gl.load(seq_lens + req).to(gl.int32) + if q_len_per_req != 1: + seq_len = seq_len - (q_len_per_req - 1) + q_offset + lens = gl.minimum(seq_len, topk).to(gl.int32) + gl.store(lens_out + row, lens) + gl.store(prefixes + row, 0) + gl.store(remaining + row, lens) + gl.store(counters + row * 2, 0) + gl.store(counters + row * 2 + 1, 0) + gl.store(out + row * out_stride + top_offsets, -1, mask=top_offsets < topk) + + +@gluon.jit +def _dsa_prefill_radix_init_kernel( + row_starts, + row_ends, + out, + lens_out, + prefixes, + remaining, + counters, + out_stride: gl.constexpr, + topk: gl.constexpr, + TOPK_LOAD_ELEMS: gl.constexpr, +): + row = gl.program_id(0) + top_layout: gl.constexpr = _vector_layout(topk, gl.num_warps(), TOPK_LOAD_ELEMS) + top_offsets = gl.arange(0, topk, layout=top_layout) + row_start = gl.load(row_starts + row).to(gl.int32) + row_end = gl.load(row_ends + row).to(gl.int32) + candidate_len = gl.maximum(row_end - row_start, 0) + lens = gl.minimum(candidate_len, topk).to(gl.int32) + gl.store(lens_out + row, lens) + gl.store(prefixes + row, 0) + gl.store(remaining + row, lens) + gl.store(counters + row * 2, 0) + gl.store(counters + row * 2 + 1, 0) + gl.store(out + row * out_stride + top_offsets, -1, mask=top_offsets < topk) + + +@gluon.jit +def _dsa_radix_hist_kernel( + logits, + prefixes, + hist, + logits_stride: gl.constexpr, + hist_tiles: gl.constexpr, + n_cols: gl.constexpr, + shift: gl.constexpr, + BLOCK_N: gl.constexpr, + LOAD_ELEMS: gl.constexpr, +): + row = gl.program_id(0) + tile = gl.program_id(1) + layout: gl.constexpr = _vector_layout(BLOCK_N, gl.num_warps(), LOAD_ELEMS) + offsets = tile * BLOCK_N + gl.arange(0, BLOCK_N, layout=layout) + mask = offsets < n_cols + values = gl.load( + logits + row * logits_stride + offsets, + mask=mask, + other=-float("inf"), + ) + keys = _fp32_to_ordered_key(values) + prefix = gl.load(prefixes + row).to(gl.uint32) + if shift == 28: + prefix_match = mask + else: + prefix_match = (keys >> (shift + 4)) == prefix + bucket = (keys >> shift) & 0xF + base = (row * hist_tiles + tile) * 16 + for bucket_id in gl.static_range(0, 16): + count = gl.sum( + gl.where( + mask & prefix_match & (bucket == bucket_id), + gl.full([BLOCK_N], 1, gl.int32, layout=layout), + gl.full([BLOCK_N], 0, gl.int32, layout=layout), + ), + axis=0, + ).to(gl.int32) + gl.store(hist + base + bucket_id, count) + + +@gluon.jit +def _dsa_radix_update_kernel( + prefixes, + remaining, + hist, + hist_tiles: gl.constexpr, + BLOCK_TILES: gl.constexpr, + LOAD_ELEMS: gl.constexpr, +): + row = gl.program_id(0) + layout: gl.constexpr = _vector_layout(BLOCK_TILES, gl.num_warps(), LOAD_ELEMS) + tile_offsets = gl.arange(0, BLOCK_TILES, layout=layout) + tile_mask = tile_offsets < hist_tiles + row_hist = hist + row * hist_tiles * 16 + kth = gl.load(remaining + row).to(gl.int32) + cumulative = 0 + selected = 0 + selected_remaining = kth + found = 0 + + for bucket_desc in gl.static_range(0, 16): + bucket_id = 15 - bucket_desc + counts = gl.load( + row_hist + tile_offsets * 16 + bucket_id, + mask=tile_mask, + other=0, + ) + count = gl.sum(counts, axis=0).to(gl.int32) + take = (found == 0) & (kth <= cumulative + count) + selected = gl.where(take, bucket_id, selected) + selected_remaining = gl.where(take, kth - cumulative, selected_remaining) + cumulative += gl.where(found == 0, count, 0) + found = gl.where(take, 1, found) + + prefix = gl.load(prefixes + row).to(gl.uint32) + gl.store(prefixes + row, ((prefix << 4) | selected).to(gl.int32)) + gl.store(remaining + row, selected_remaining) + + +@gluon.jit +def _dsa_decode_radix_scatter_slots_kernel( + logits, + prefixes, + remaining, + counters, + block_table, + seq_lens, + out, + logits_stride: gl.constexpr, + block_table_stride: gl.constexpr, + out_stride: gl.constexpr, + block_table_cols: gl.constexpr, + n_cols: gl.constexpr, + page_size: gl.constexpr, + topk: gl.constexpr, + q_len_per_req: gl.constexpr, + BLOCK_N: gl.constexpr, + LOAD_ELEMS: gl.constexpr, +): + row = gl.program_id(0) + tile = gl.program_id(1) + layout: gl.constexpr = _vector_layout(BLOCK_N, gl.num_warps(), LOAD_ELEMS) + offsets = tile * BLOCK_N + gl.arange(0, BLOCK_N, layout=layout) + req = row // q_len_per_req + q_offset = row - req * q_len_per_req + seq_len = gl.load(seq_lens + req).to(gl.int32) + if q_len_per_req != 1: + seq_len = seq_len - (q_len_per_req - 1) + q_offset + mask = (offsets < n_cols) & (offsets < seq_len) + values = gl.load( + logits + row * logits_stride + offsets, + mask=mask, + other=-float("inf"), + ) + threshold = gl.load(prefixes + row).to(gl.uint32) + keep_equal = gl.load(remaining + row).to(gl.int32) + count_greater = topk - keep_equal + finite = values != -float("inf") + keys = _fp32_to_ordered_key(values) + greater = mask & finite & (keys > threshold) + equal = mask & finite & (keys == threshold) + + greater_i32 = greater.to(gl.int32) + equal_i32 = equal.to(gl.int32) + tile_greater = gl.sum(greater_i32, axis=0).to(gl.int32) + tile_equal = gl.sum(equal_i32, axis=0).to(gl.int32) + greater_start = gl.atomic_add( + counters + row * 2, tile_greater, sem="acq_rel", scope="gpu" + ) + equal_start = gl.atomic_add( + counters + row * 2 + 1, tile_equal, sem="acq_rel", scope="gpu" + ) + + greater_pos = greater_start + gl.associative_scan(greater_i32, 0, _topk_add) - 1 + equal_pos = ( + count_greater + equal_start + gl.associative_scan(equal_i32, 0, _topk_add) - 1 + ) + greater_write = greater & (greater_pos < topk) + equal_write = equal & (equal_pos < topk) & (equal_pos < count_greater + keep_equal) + + block_idx = offsets.to(gl.int32) // page_size + block_offset = offsets.to(gl.int32) - block_idx * page_size + page = gl.load( + block_table + req * block_table_stride + block_idx, + mask=(greater_write | equal_write) & (block_idx < block_table_cols), + other=0, + ).to(gl.int32) + slots = page * page_size + block_offset + gl.store(out + row * out_stride + greater_pos, slots, mask=greater_write) + gl.store(out + row * out_stride + equal_pos, slots, mask=equal_write) + + +@gluon.jit +def _dsa_prefill_radix_scatter_kernel( + logits, + prefixes, + remaining, + counters, + row_starts, + row_ends, + out, + logits_stride: gl.constexpr, + out_stride: gl.constexpr, + n_cols: gl.constexpr, + topk: gl.constexpr, + BLOCK_N: gl.constexpr, + LOAD_ELEMS: gl.constexpr, +): + row = gl.program_id(0) + tile = gl.program_id(1) + layout: gl.constexpr = _vector_layout(BLOCK_N, gl.num_warps(), LOAD_ELEMS) + offsets = tile * BLOCK_N + gl.arange(0, BLOCK_N, layout=layout) + row_start = gl.load(row_starts + row).to(gl.int32) + row_end = gl.load(row_ends + row).to(gl.int32) + mask = (offsets >= row_start) & (offsets < row_end) & (offsets < n_cols) + values = gl.load( + logits + row * logits_stride + offsets, + mask=mask, + other=-float("inf"), + ) + threshold = gl.load(prefixes + row).to(gl.uint32) + keep_equal = gl.load(remaining + row).to(gl.int32) + count_greater = topk - keep_equal + finite = values != -float("inf") + keys = _fp32_to_ordered_key(values) + greater = mask & finite & (keys > threshold) + equal = mask & finite & (keys == threshold) + + greater_i32 = greater.to(gl.int32) + equal_i32 = equal.to(gl.int32) + tile_greater = gl.sum(greater_i32, axis=0).to(gl.int32) + tile_equal = gl.sum(equal_i32, axis=0).to(gl.int32) + greater_start = gl.atomic_add( + counters + row * 2, tile_greater, sem="acq_rel", scope="gpu" + ) + equal_start = gl.atomic_add( + counters + row * 2 + 1, tile_equal, sem="acq_rel", scope="gpu" + ) + + greater_pos = greater_start + gl.associative_scan(greater_i32, 0, _topk_add) - 1 + equal_pos = ( + count_greater + equal_start + gl.associative_scan(equal_i32, 0, _topk_add) - 1 + ) + greater_write = greater & (greater_pos < topk) + equal_write = equal & (equal_pos < topk) & (equal_pos < count_greater + keep_equal) + local = offsets.to(gl.int32) + gl.store(out + row * out_stride + greater_pos, local, mask=greater_write) + gl.store(out + row * out_stride + equal_pos, local, mask=equal_write) + + +def _load_elems(block: int, num_warps: int) -> int: + return max(1, triton.cdiv(int(block), 64 * int(num_warps))) + + +def _contiguous(tensor: torch.Tensor) -> torch.Tensor: + return tensor if tensor.is_contiguous() else tensor.contiguous() + + +def _to_contiguous( + tensor: torch.Tensor, + *, + device: torch.device, + dtype: torch.dtype, +) -> torch.Tensor: + tensor = tensor.to(device=device, dtype=dtype) + return _contiguous(tensor) + + +def _validate_topk(topk: int) -> None: + if topk <= 0: + raise ValueError(f"topk must be positive, got {topk}") + if topk & (topk - 1): + raise ValueError(f"DSA Gluon top-k requires power-of-two topk, got {topk}") + + +def _use_radix_topk(cols: int) -> bool: + return int(cols) >= _RADIX_TOPK_MIN_COLS + + +def _dsa_radix_scratch( + rows: int, + cols: int, + *, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + tiles = triton.cdiv(int(cols), _RADIX_TOPK_BLOCK_N) + hist = torch.empty((rows, tiles, 16), dtype=torch.int32, device=device) + prefixes = torch.empty((rows,), dtype=torch.int32, device=device) + remaining = torch.empty((rows,), dtype=torch.int32, device=device) + counters = torch.empty((rows, 2), dtype=torch.int32, device=device) + block_tiles = triton.next_power_of_2(tiles) + return hist, prefixes, remaining, counters, tiles, block_tiles + + +def _run_radix_prefix_passes( + logits: torch.Tensor, + hist: torch.Tensor, + prefixes: torch.Tensor, + remaining: torch.Tensor, + *, + rows: int, + cols: int, + tiles: int, + block_tiles: int, +) -> None: + hist_load_elems = _load_elems(_RADIX_TOPK_BLOCK_N, 8) + update_load_elems = _load_elems(block_tiles, 8) + # Ordered FP32 keys have 8 nibbles; each pass fixes one more prefix nibble. + for shift in range(28, -1, -4): + _dsa_radix_hist_kernel[(rows, tiles)]( + logits, + prefixes, + hist, + logits.stride(0), + tiles, + n_cols=cols, + shift=shift, + BLOCK_N=_RADIX_TOPK_BLOCK_N, + LOAD_ELEMS=hist_load_elems, + num_warps=8, + ) + _dsa_radix_update_kernel[(rows,)]( + prefixes, + remaining, + hist, + hist_tiles=tiles, + BLOCK_TILES=block_tiles, + LOAD_ELEMS=update_load_elems, + num_warps=8, + ) + + +def _dsa_decode_radix_topk_slots( + logits: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + *, + page_size: int, + topk: int, + q_len_per_req: int, + out: torch.Tensor, + lens_out: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + rows, cols = logits.shape + hist, prefixes, remaining, counters, tiles, block_tiles = _dsa_radix_scratch( + rows, cols, device=logits.device + ) + _dsa_decode_radix_init_kernel[(rows,)]( + seq_lens, + out, + lens_out, + prefixes, + remaining, + counters, + out.stride(0), + topk=topk, + q_len_per_req=q_len_per_req, + TOPK_LOAD_ELEMS=_load_elems(topk, 8), + num_warps=8, + ) + _run_radix_prefix_passes( + logits, + hist, + prefixes, + remaining, + rows=rows, + cols=cols, + tiles=tiles, + block_tiles=block_tiles, + ) + _dsa_decode_radix_scatter_slots_kernel[(rows, tiles)]( + logits, + prefixes, + remaining, + counters, + block_table, + seq_lens, + out, + logits.stride(0), + block_table.stride(0), + out.stride(0), + block_table.shape[1], + n_cols=cols, + page_size=int(page_size), + topk=topk, + q_len_per_req=q_len_per_req, + BLOCK_N=_RADIX_TOPK_BLOCK_N, + LOAD_ELEMS=_load_elems(_RADIX_TOPK_BLOCK_N, 8), + num_warps=8, + ) + return out, lens_out + + +def _dsa_prefill_radix_topk( + logits: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + *, + topk: int, + out: torch.Tensor, + lens_out: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + rows, cols = logits.shape + hist, prefixes, remaining, counters, tiles, block_tiles = _dsa_radix_scratch( + rows, cols, device=logits.device + ) + _dsa_prefill_radix_init_kernel[(rows,)]( + row_starts, + row_ends, + out, + lens_out, + prefixes, + remaining, + counters, + out.stride(0), + topk=topk, + TOPK_LOAD_ELEMS=_load_elems(topk, 8), + num_warps=8, + ) + _run_radix_prefix_passes( + logits, + hist, + prefixes, + remaining, + rows=rows, + cols=cols, + tiles=tiles, + block_tiles=block_tiles, + ) + _dsa_prefill_radix_scatter_kernel[(rows, tiles)]( + logits, + prefixes, + remaining, + counters, + row_starts, + row_ends, + out, + logits.stride(0), + out.stride(0), + n_cols=cols, + topk=topk, + BLOCK_N=_RADIX_TOPK_BLOCK_N, + LOAD_ELEMS=_load_elems(_RADIX_TOPK_BLOCK_N, 8), + num_warps=8, + ) + return out, lens_out + + +def gluon_dsa_decode_topk_fp8_gfx950( + q: torch.Tensor, + weights: torch.Tensor, + seq_lens: torch.Tensor, + block_table: torch.Tensor, + *, + page_size: int, + topk: int, + softmax_scale: float, + q_len_per_req: int = 1, + index_k_cache: torch.Tensor | None = None, + seq_lens_2d: torch.Tensor | None = None, + plan: object | None = None, + out: torch.Tensor | None = None, + lens_out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + del plan, seq_lens_2d + topk = int(topk) + q_len_per_req = int(q_len_per_req) + _validate_topk(topk) + if not 1 <= q_len_per_req <= 6: + raise ValueError(f"q_len_per_req must be in [1, 6], got {q_len_per_req}") + if index_k_cache is None: + raise RuntimeError("Gluon DSA paged top-k requires packed FP8 index_k_cache") + row_bytes = _check_packed_fp8_inputs(q, index_k_cache, weights, int(page_size)) + if seq_lens.dim() != 1: + raise ValueError( + f"seq_lens must be 1-D, got {tuple(seq_lens.shape)} for q={tuple(q.shape)}" + ) + expected_tokens = int(seq_lens.numel()) * q_len_per_req + if expected_tokens != q.shape[0]: + raise ValueError( + "q rows must equal seq_lens rows times q_len_per_req, got " + f"q={tuple(q.shape)}, seq_lens={tuple(seq_lens.shape)}, " + f"q_len_per_req={q_len_per_req}" + ) + if block_table.dim() != 2 or block_table.shape[0] < seq_lens.numel(): + raise ValueError( + "block_table must have at least one row per request, got " + f"block_table={tuple(block_table.shape)}, q={tuple(q.shape)}" + ) + if q.shape[0] == 0: + empty_out = ( + torch.empty((0, int(topk)), dtype=torch.int32, device=q.device) + if out is None + else out + ) + empty_lens = ( + torch.empty((0,), dtype=torch.int32, device=q.device) + if lens_out is None + else lens_out + ) + return empty_out, empty_lens + q = _contiguous(q) + index_k_cache = _contiguous(index_k_cache) + weights = _contiguous(weights) + seq_lens = _to_contiguous(seq_lens, device=q.device, dtype=torch.int32) + block_table = _to_contiguous(block_table, device=q.device, dtype=torch.int32) + max_seq_len = int(block_table.shape[1]) * int(page_size) + if out is None: + out = torch.empty((q.shape[0], topk), dtype=torch.int32, device=q.device) + if lens_out is None: + lens_out = torch.empty((q.shape[0],), dtype=torch.int32, device=q.device) + logits = torch.empty( + (q.shape[0], max_seq_len), dtype=torch.float32, device=q.device + ) + block_n = 32 + _dsa_decode_logits_fp8_kernel[(q.shape[0], triton.cdiv(max_seq_len, block_n))]( + q, + index_k_cache.view(torch.float8_e4m3fn), + index_k_cache.view(torch.float32), + weights, + seq_lens, + block_table, + logits, + block_table.stride(0), + logits.stride(0), + page_size=int(page_size), + row_bytes=row_bytes, + max_seq_len=max_seq_len, + num_heads=q.shape[1], + head_dim=q.shape[2], + num_groups=q.shape[2] // 128, + softmax_scale=float(softmax_scale), + q_len_per_req=q_len_per_req, + BLOCK_N=block_n, + BLOCK_D=128, + num_warps=4, + ) + if _use_radix_topk(max_seq_len): + return _dsa_decode_radix_topk_slots( + logits, + block_table, + seq_lens, + page_size=int(page_size), + topk=topk, + q_len_per_req=q_len_per_req, + out=out, + lens_out=lens_out, + ) + + select_warps = 8 + select_block = triton.next_power_of_2(max(max_seq_len, topk)) + _dsa_decode_select_topk_kernel[(q.shape[0],)]( + logits, + block_table, + seq_lens, + out, + lens_out, + logits.stride(0), + block_table.stride(0), + out.stride(0), + block_table.shape[1], + page_size=int(page_size), + topk=topk, + q_len_per_req=q_len_per_req, + BLOCK_N=select_block, + LOAD_ELEMS=_load_elems(select_block, select_warps), + TOPK_LOAD_ELEMS=_load_elems(topk, select_warps), + num_warps=select_warps, + ) + return out, lens_out + + +def gluon_dsa_prefill_topk_fp8_gfx950( + q: torch.Tensor, + weights: torch.Tensor, + kv_workspace_slots: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + *, + topk: int, + softmax_scale: float, + index_k_cache: torch.Tensor | None = None, + page_size: int | None = None, + index_k_fp8: torch.Tensor | None = None, + index_k_scale: torch.Tensor | None = None, + max_logits_bytes: int | None = None, + out: torch.Tensor | None = None, + lens_out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + del index_k_fp8, index_k_scale + topk = int(topk) + _validate_topk(topk) + if index_k_cache is None or page_size is None: + raise RuntimeError( + "Gluon DSA top-k requires packed FP8 index_k_cache and page_size" + ) + row_bytes = _check_packed_fp8_inputs(q, index_k_cache, weights, int(page_size)) + if kv_workspace_slots.dim() != 1: + raise ValueError( + f"kv_workspace_slots must be 1-D, got {tuple(kv_workspace_slots.shape)}" + ) + if row_starts.shape != (q.shape[0],) or row_ends.shape != (q.shape[0],): + raise ValueError( + "row_starts/row_ends must be [tokens], got " + f"row_starts={tuple(row_starts.shape)}, row_ends={tuple(row_ends.shape)}, " + f"q={tuple(q.shape)}" + ) + if out is None: + out = torch.empty((q.shape[0], topk), dtype=torch.int32, device=q.device) + if lens_out is None: + lens_out = torch.empty((q.shape[0],), dtype=torch.int32, device=q.device) + if q.shape[0] == 0: + return out, lens_out + q = _contiguous(q) + index_k_cache = _contiguous(index_k_cache) + weights = _contiguous(weights) + kv_workspace_slots = _to_contiguous( + kv_workspace_slots, device=q.device, dtype=torch.int64 + ) + row_starts = _to_contiguous(row_starts, device=q.device, dtype=torch.int32) + row_ends = _to_contiguous(row_ends, device=q.device, dtype=torch.int32) + seq_len_sum = int(kv_workspace_slots.numel()) + if seq_len_sum == 0: + out.fill_(-1) + lens_out.zero_() + return out, lens_out + + if max_logits_bytes is None: + max_query_rows = q.shape[0] + else: + max_query_rows = max(1, int(max_logits_bytes) // (max(seq_len_sum, 1) * 4)) + block_n = 32 + select_warps = 8 + select_block = triton.next_power_of_2(max(seq_len_sum, topk)) + for start in range(0, q.shape[0], max_query_rows): + end = min(start + max_query_rows, q.shape[0]) + logits = torch.empty( + (end - start, seq_len_sum), dtype=torch.float32, device=q.device + ) + _dsa_prefill_logits_fp8_kernel[ + (end - start, triton.cdiv(seq_len_sum, block_n)) + ]( + q[start:end], + index_k_cache.view(torch.float8_e4m3fn), + index_k_cache.view(torch.float32), + weights[start:end], + kv_workspace_slots, + row_starts[start:end], + row_ends[start:end], + logits, + logits.stride(0), + seq_len_sum=seq_len_sum, + page_size=int(page_size), + row_bytes=row_bytes, + num_heads=q.shape[1], + head_dim=q.shape[2], + num_groups=q.shape[2] // 128, + softmax_scale=float(softmax_scale), + BLOCK_N=block_n, + BLOCK_D=128, + num_warps=4, + ) + if _use_radix_topk(seq_len_sum): + _dsa_prefill_radix_topk( + logits, + row_starts[start:end], + row_ends[start:end], + topk=topk, + out=out[start:end], + lens_out=lens_out[start:end], + ) + continue + + _dsa_prefill_select_topk_kernel[(end - start,)]( + logits, + row_starts[start:end], + row_ends[start:end], + out[start:end], + lens_out[start:end], + logits.stride(0), + out.stride(0), + topk=topk, + BLOCK_N=select_block, + LOAD_ELEMS=_load_elems(select_block, select_warps), + TOPK_LOAD_ELEMS=_load_elems(topk, select_warps), + num_warps=select_warps, + ) + return out, lens_out diff --git a/tokenspeed-kernel-amd/test/ops/test_attention_dsa_gluon_gfx950.py b/tokenspeed-kernel-amd/test/ops/test_attention_dsa_gluon_gfx950.py new file mode 100644 index 000000000..fd9dd723c --- /dev/null +++ b/tokenspeed-kernel-amd/test/ops/test_attention_dsa_gluon_gfx950.py @@ -0,0 +1,1273 @@ +# Copyright (c) 2026 LightSeek Foundation + +from __future__ import annotations + +import math +from collections.abc import Sequence +from dataclasses import dataclass + +import pytest +import torch + + +def _is_gfx950() -> bool: + if not torch.cuda.is_available(): + return False + arch = getattr(torch.cuda.get_device_properties(0), "gcnArchName", "") + return "gfx950" in arch + + +if not _is_gfx950(): + pytest.skip("AMD GFX950 is required for Gluon DSA tests", allow_module_level=True) + + +from tokenspeed_kernel_amd.ops.attention.gluon import dsa_topk_gfx950 # noqa: E402 +from tokenspeed_kernel_amd.ops.attention.gluon.dsa_gfx950 import ( # noqa: E402 + _trim_topk_slots_for_context, + gluon_dsa_decode_gfx950, + gluon_dsa_prefill_gfx950, +) +from tokenspeed_kernel_amd.ops.attention.gluon.dsa_topk_gfx950 import ( # noqa: E402 + gluon_dsa_decode_topk_fp8_gfx950, + gluon_dsa_prefill_topk_fp8_gfx950, +) + +torch.manual_seed(42) + + +@pytest.mark.parametrize( + ("max_seqlen_k", "expected_topk"), + ((25, 512), (608, 1024), (1537, 2048)), +) +def test_dsa_attention_trims_topk_to_registered_context_width( + max_seqlen_k: int, + expected_topk: int, +) -> None: + topk_slots = torch.arange(2 * 2048, device="cuda", dtype=torch.int32).reshape( + 2, 2048 + ) + + trimmed = _trim_topk_slots_for_context(topk_slots, max_seqlen_k) + + assert trimmed.shape == (2, expected_topk) + torch.testing.assert_close(trimmed, topk_slots[:, :expected_topk]) + + +@dataclass(frozen=True) +class _TopKDecodeCase: + name: str + seq_lens: tuple[int, ...] + index_heads: int + topk: int + q_len_per_req: int + seed: int + + +@dataclass(frozen=True) +class _TopKPrefillCase: + name: str + prefix_lens: tuple[int, ...] + extend_lens: tuple[int, ...] + index_heads: int + topk: int + seed: int + + +@dataclass(frozen=True) +class _DSACase: + name: str + mode: str + kv_layout: str + topk: int + seed: int + num_heads: int = 8 + qk_nope_head_dim: int = 192 + kv_lora_rank: int = 512 + qk_rope_head_dim: int = 64 + q_len_per_req: int = 1 + visible_lens: tuple[int, ...] | None = None + topk_lens: tuple[int, ...] | None = None + prefix_lens: tuple[int, ...] | None = None + extend_lens: tuple[int, ...] | None = None + + +_GLM52_TOPK_DECODE_CASES = ( + _TopKDecodeCase( + "decode_batch_mixed_512", + seq_lens=(128, 257, 511, 1024), + index_heads=2, + topk=512, + q_len_per_req=1, + seed=101, + ), + _TopKDecodeCase( + "decode_q3_boundary_512", + seq_lens=(510, 511, 512, 1022, 1023, 1024), + index_heads=2, + topk=512, + q_len_per_req=3, + seed=102, + ), + _TopKDecodeCase( + "decode_long_1024", + seq_lens=(2048, 3072, 4096), + index_heads=4, + topk=1024, + q_len_per_req=1, + seed=103, + ), + _TopKDecodeCase( + "decode_long_2048", + seq_lens=(1536, 4096), + index_heads=2, + topk=2048, + q_len_per_req=1, + seed=104, + ), +) + + +_GLM52_TOPK_PREFILL_CASES = ( + _TopKPrefillCase( + "prefill_short_512", + prefix_lens=(64, 128), + extend_lens=(16, 32), + index_heads=2, + topk=512, + seed=201, + ), + _TopKPrefillCase( + "prefill_chunk_512", + prefix_lens=(512, 1024), + extend_lens=(32, 32), + index_heads=2, + topk=512, + seed=202, + ), + _TopKPrefillCase( + "prefill_mixed_1024", + prefix_lens=(256, 1024, 1536), + extend_lens=(16, 24, 16), + index_heads=4, + topk=1024, + seed=203, + ), + _TopKPrefillCase( + "prefill_long_2048", + prefix_lens=(1536, 2048), + extend_lens=(16, 16), + index_heads=2, + topk=2048, + seed=204, + ), +) + + +_GLM52_DSA_CASES = ( + _DSACase( + "decode_sparse_mixed_512", + mode="decode", + kv_layout="sparse", + topk=512, + visible_lens=(128, 257, 512, 1024), + topk_lens=(64, 257, 512, 384), + seed=301, + ), + _DSACase( + "decode_dense_q3_512", + mode="decode", + kv_layout="dense", + topk=512, + q_len_per_req=3, + visible_lens=(512, 513, 514, 1024, 1025, 1026), + topk_lens=(128, 256, 512, 300, 511, 64), + seed=302, + ), + _DSACase( + "decode_sparse_long_1024", + mode="decode", + kv_layout="sparse", + topk=1024, + visible_lens=(2048, 3072, 4096), + topk_lens=(640, 1024, 777), + seed=303, + ), + _DSACase( + "decode_dense_long_2048", + mode="decode", + kv_layout="dense", + topk=2048, + visible_lens=(2048, 4096), + topk_lens=(1536, 2048), + seed=304, + ), + _DSACase( + "prefill_sparse_short_512", + mode="prefill", + kv_layout="sparse", + topk=512, + prefix_lens=(64, 128), + extend_lens=(8, 8), + topk_lens=( + 32, + 64, + 96, + 128, + 48, + 80, + 112, + 136, + 33, + 65, + 97, + 129, + 49, + 81, + 113, + 136, + ), + seed=305, + ), + _DSACase( + "prefill_dense_chunk_512", + mode="prefill", + kv_layout="dense", + topk=512, + prefix_lens=(512, 1024), + extend_lens=(8, 8), + topk_lens=( + 128, + 192, + 256, + 320, + 384, + 448, + 512, + 256, + 96, + 160, + 224, + 288, + 352, + 416, + 480, + 512, + ), + seed=306, + ), + _DSACase( + "prefill_sparse_mixed_1024", + mode="prefill", + kv_layout="sparse", + topk=1024, + prefix_lens=(256, 1024, 1536), + extend_lens=(4, 6, 4), + topk_lens=( + 128, + 256, + 384, + 260, + 512, + 640, + 768, + 896, + 1024, + 768, + 512, + 1024, + 768, + 1024, + ), + seed=307, + ), + _DSACase( + "prefill_dense_long_2048", + mode="prefill", + kv_layout="dense", + topk=2048, + prefix_lens=(1536, 2048), + extend_lens=(4, 4), + topk_lens=(512, 1024, 1536, 1539, 1024, 1536, 2048, 2048), + seed=308, + ), + _DSACase( + "prefill_sparse_first_prompt_2048", + mode="prefill", + kv_layout="sparse", + topk=2048, + num_heads=16, + prefix_lens=(0,), + extend_lens=(25,), + topk_lens=tuple(range(1, 26)), + seed=309, + ), +) + + +def _pack_index_k_cache( + index_k: torch.Tensor, + page_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + head_dim = index_k.shape[1] + num_groups = head_dim // 128 + row_bytes = head_dim + num_groups * 4 + num_slots = index_k.shape[0] + num_pages = num_slots // page_size + packed = torch.empty( + (num_slots, row_bytes), + device=index_k.device, + dtype=torch.uint8, + ) + x = index_k.float().reshape(num_slots, num_groups, 128) + scale = x.abs().amax(dim=-1, keepdim=True).clamp_min(1.0e-6) / 448.0 + x_fp8 = (x / scale).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + + flat = packed.reshape(-1) + page_bytes = page_size * row_bytes + fp8_view = torch.as_strided( + flat.view(torch.float8_e4m3fn), + (num_pages, page_size, head_dim), + (page_bytes, head_dim, 1), + ) + scale_view = torch.as_strided( + flat.view(torch.float32), + (num_pages, page_size, num_groups), + (page_bytes // 4, num_groups, 1), + (page_size * head_dim) // 4, + ) + fp8_view.copy_(x_fp8.reshape(num_pages, page_size, head_dim)) + scale_view.copy_(scale.reshape(num_pages, page_size, num_groups)) + return packed, (x_fp8.float() * scale).reshape_as(index_k) + + +def _generator(device: str, seed: int) -> torch.Generator: + gen = torch.Generator(device=device) + gen.manual_seed(seed) + return gen + + +def _randn_bf16( + shape: Sequence[int], + *, + device: str, + generator: torch.Generator, + scale: float = 0.25, +) -> torch.Tensor: + return ( + torch.randn(shape, device=device, dtype=torch.float32, generator=generator) + * scale + ).to(torch.bfloat16) + + +def _normal_weights( + shape: Sequence[int], + *, + device: str, + generator: torch.Generator, +) -> torch.Tensor: + logits = torch.randn(shape, device=device, dtype=torch.float32, generator=generator) + return torch.softmax(logits, dim=-1).contiguous() + + +def _round_up_to_page(slots: int, page_size: int) -> int: + return int(math.ceil(slots / page_size) * page_size) + + +def _make_decode_block_table( + seq_lens: Sequence[int], + page_size: int, + device: str, +) -> tuple[torch.Tensor, int]: + max_pages = max(math.ceil(seq_len / page_size) for seq_len in seq_lens) + pages = torch.arange( + len(seq_lens) * max_pages, device=device, dtype=torch.int32 + ).reshape(len(seq_lens), max_pages) + return pages, int(len(seq_lens) * max_pages * page_size) + + +def _make_prefill_workspace( + prefix_lens: Sequence[int], + extend_lens: Sequence[int], + *, + device: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[range]]: + kv_workspace_slots: list[int] = [] + row_starts: list[int] = [] + row_ends: list[int] = [] + visible_ranges: list[range] = [] + cursor = 0 + for prefix_len, extend_len in zip(prefix_lens, extend_lens, strict=True): + req_start = cursor + seq_len = int(prefix_len) + int(extend_len) + kv_workspace_slots.extend(range(req_start, req_start + seq_len)) + for query_offset in range(int(extend_len)): + visible_end = req_start + int(prefix_len) + query_offset + 1 + row_starts.append(req_start) + row_ends.append(visible_end) + visible_ranges.append(range(req_start, visible_end)) + cursor += seq_len + + return ( + torch.tensor(kv_workspace_slots, device=device, dtype=torch.int64), + torch.tensor(row_starts, device=device, dtype=torch.int32), + torch.tensor(row_ends, device=device, dtype=torch.int32), + visible_ranges, + ) + + +def _index_scores( + q: torch.Tensor, + weights: torch.Tensor, + index_k: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + per_head = index_k.float() @ q.float().transpose(0, 1) + return (per_head * weights.float()).sum(dim=1) * softmax_scale + + +def _reference_decode_topk( + q: torch.Tensor, + weights: torch.Tensor, + index_k: torch.Tensor, + seq_lens: torch.Tensor, + block_table: torch.Tensor, + *, + page_size: int, + topk: int, + softmax_scale: float, + q_len_per_req: int = 1, +) -> tuple[torch.Tensor, torch.Tensor]: + out = torch.full((q.shape[0], topk), -1, device=q.device, dtype=torch.int32) + lens = torch.empty((q.shape[0],), device=q.device, dtype=torch.int32) + for token in range(q.shape[0]): + req = token // int(q_len_per_req) + q_offset = token - req * int(q_len_per_req) + seq_len = int(seq_lens[req].item()) + if q_len_per_req != 1: + seq_len = seq_len - (int(q_len_per_req) - 1) + q_offset + count = min(seq_len, int(topk)) + lens[token] = count + if count == 0: + continue + offsets = torch.arange(seq_len, device=q.device, dtype=torch.long) + pages = block_table[req].long().index_select(0, offsets // page_size) + slots = pages * int(page_size) + offsets.remainder(page_size) + scores = _index_scores( + q[token], + weights[token], + index_k.index_select(0, slots), + softmax_scale, + ) + selected = torch.topk(scores, count).indices + out[token, :count] = slots.index_select(0, selected).to(torch.int32) + return out, lens + + +def _reference_prefill_topk( + q: torch.Tensor, + weights: torch.Tensor, + index_k: torch.Tensor, + kv_workspace_slots: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + *, + topk: int, + softmax_scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + out = torch.full((q.shape[0], topk), -1, device=q.device, dtype=torch.int32) + candidate_lens = (row_ends - row_starts).clamp_min(0) + lens = torch.minimum(candidate_lens, torch.full_like(candidate_lens, int(topk))) + for token in range(q.shape[0]): + count = int(lens[token].item()) + if count == 0: + continue + rows = torch.arange( + int(row_starts[token].item()), + int(row_ends[token].item()), + device=q.device, + dtype=torch.long, + ) + slots = kv_workspace_slots.index_select(0, rows).long() + scores = _index_scores( + q[token], + weights[token], + index_k.index_select(0, slots), + softmax_scale, + ) + selected = torch.topk(scores, count).indices + out[token, :count] = rows.index_select(0, selected).to(torch.int32) + return out, lens + + +def _assert_topk_matches( + actual: torch.Tensor, + actual_lens: torch.Tensor, + expected: torch.Tensor, + expected_lens: torch.Tensor, +) -> None: + torch.testing.assert_close(actual_lens.cpu(), expected_lens.cpu()) + for token in range(actual.shape[0]): + count = int(expected_lens[token].item()) + actual_selected = torch.sort(actual[token, :count].cpu()).values + expected_selected = torch.sort(expected[token, :count].cpu()).values + torch.testing.assert_close(actual_selected, expected_selected) + assert (actual[token, count:] == -1).all() + + +def _strided_last_dim(tensor: torch.Tensor) -> torch.Tensor: + backing = torch.empty( + (*tensor.shape[:-1], tensor.shape[-1] * 2), + device=tensor.device, + dtype=tensor.dtype, + ) + view = backing[..., ::2] + view.copy_(tensor) + return view + + +def _strided_1d(tensor: torch.Tensor) -> torch.Tensor: + backing = torch.empty( + (tensor.shape[0] * 2,), + device=tensor.device, + dtype=tensor.dtype, + ) + view = backing[::2] + view.copy_(tensor) + return view + + +@pytest.mark.parametrize( + "case", + _GLM52_TOPK_DECODE_CASES, + ids=lambda case: case.name, +) +def test_dsa_decode_topk_fp8_glm52_cases(case: _TopKDecodeCase) -> None: + device = "cuda" + page_size = 64 + head_dim = 128 + softmax_scale = head_dim**-0.5 + gen = _generator(device, case.seed) + block_table, num_slots = _make_decode_block_table(case.seq_lens, page_size, device) + tokens = len(case.seq_lens) * case.q_len_per_req + q = _randn_bf16( + (tokens, case.index_heads, head_dim), + device=device, + generator=gen, + ) + weights = _normal_weights((tokens, case.index_heads), device=device, generator=gen) + packed_index_k, index_k = _pack_index_k_cache( + _randn_bf16((num_slots, head_dim), device=device, generator=gen), + page_size, + ) + seq_lens = torch.tensor(case.seq_lens, device=device, dtype=torch.int32) + + topk_slots, topk_lens = gluon_dsa_decode_topk_fp8_gfx950( + q, + weights, + seq_lens, + block_table, + page_size=page_size, + topk=case.topk, + softmax_scale=softmax_scale, + seq_lens_2d=seq_lens.unsqueeze(1).expand(-1, case.q_len_per_req), + q_len_per_req=case.q_len_per_req, + index_k_cache=packed_index_k, + ) + expected_slots, expected_lens = _reference_decode_topk( + q, + weights, + index_k, + seq_lens, + block_table, + page_size=page_size, + topk=case.topk, + softmax_scale=softmax_scale, + q_len_per_req=case.q_len_per_req, + ) + + _assert_topk_matches(topk_slots, topk_lens, expected_slots, expected_lens) + + +def test_dsa_decode_topk_fp8_accepts_strided_inputs() -> None: + device = "cuda" + page_size = 64 + head_dim = 128 + topk = 512 + softmax_scale = head_dim**-0.5 + gen = _generator(device, 121) + seq_lens_tuple = (640, 704) + block_table, num_slots = _make_decode_block_table(seq_lens_tuple, page_size, device) + tokens = len(seq_lens_tuple) + q = _strided_last_dim( + _randn_bf16((tokens, 1, head_dim), device=device, generator=gen) + ) + weights = _strided_last_dim( + _normal_weights((tokens, 1), device=device, generator=gen) + ) + packed_index_k, index_k = _pack_index_k_cache( + _randn_bf16((num_slots, head_dim), device=device, generator=gen), + page_size, + ) + seq_lens = _strided_1d( + torch.tensor(seq_lens_tuple, device=device, dtype=torch.int32) + ) + block_table = _strided_last_dim(block_table) + packed_index_k = _strided_last_dim(packed_index_k) + + topk_slots, topk_lens = gluon_dsa_decode_topk_fp8_gfx950( + q, + weights, + seq_lens, + block_table, + page_size=page_size, + topk=topk, + softmax_scale=softmax_scale, + q_len_per_req=1, + index_k_cache=packed_index_k, + ) + expected_slots, expected_lens = _reference_decode_topk( + q, + weights, + index_k, + seq_lens, + block_table, + page_size=page_size, + topk=topk, + softmax_scale=softmax_scale, + ) + + _assert_topk_matches(topk_slots, topk_lens, expected_slots, expected_lens) + + +@pytest.mark.parametrize( + "case", + _GLM52_TOPK_PREFILL_CASES, + ids=lambda case: case.name, +) +def test_dsa_prefill_topk_fp8_glm52_cases(case: _TopKPrefillCase) -> None: + device = "cuda" + page_size = 64 + head_dim = 128 + softmax_scale = head_dim**-0.5 + gen = _generator(device, case.seed) + kv_workspace_slots, row_starts, row_ends, _ = _make_prefill_workspace( + case.prefix_lens, case.extend_lens, device=device + ) + num_tokens = int(sum(case.extend_lens)) + num_slots = _round_up_to_page(int(kv_workspace_slots.numel()), page_size) + q = _randn_bf16( + (num_tokens, case.index_heads, head_dim), + device=device, + generator=gen, + ) + weights = _normal_weights( + (num_tokens, case.index_heads), device=device, generator=gen + ) + packed_index_k, index_k = _pack_index_k_cache( + _randn_bf16((num_slots, head_dim), device=device, generator=gen), + page_size, + ) + + workspace_indices, topk_lens = gluon_dsa_prefill_topk_fp8_gfx950( + q, + weights, + kv_workspace_slots, + row_starts, + row_ends, + topk=case.topk, + softmax_scale=softmax_scale, + index_k_cache=packed_index_k, + page_size=page_size, + ) + expected_indices, expected_lens = _reference_prefill_topk( + q, + weights, + index_k, + kv_workspace_slots, + row_starts, + row_ends, + topk=case.topk, + softmax_scale=softmax_scale, + ) + + _assert_topk_matches(workspace_indices, topk_lens, expected_indices, expected_lens) + + +def test_dsa_prefill_topk_fp8_accepts_strided_inputs() -> None: + device = "cuda" + page_size = 64 + head_dim = 128 + topk = 512 + softmax_scale = head_dim**-0.5 + gen = _generator(device, 221) + kv_workspace_slots, row_starts, row_ends, _ = _make_prefill_workspace( + (640,), (2,), device=device + ) + num_tokens = int(row_starts.numel()) + num_slots = _round_up_to_page(int(kv_workspace_slots.numel()), page_size) + q = _strided_last_dim( + _randn_bf16((num_tokens, 1, head_dim), device=device, generator=gen) + ) + weights = _strided_last_dim( + _normal_weights((num_tokens, 1), device=device, generator=gen) + ) + packed_index_k, index_k = _pack_index_k_cache( + _randn_bf16((num_slots, head_dim), device=device, generator=gen), + page_size, + ) + kv_workspace_slots = _strided_1d(kv_workspace_slots) + row_starts = _strided_1d(row_starts) + row_ends = _strided_1d(row_ends) + packed_index_k = _strided_last_dim(packed_index_k) + + workspace_indices, topk_lens = gluon_dsa_prefill_topk_fp8_gfx950( + q, + weights, + kv_workspace_slots, + row_starts, + row_ends, + topk=topk, + softmax_scale=softmax_scale, + index_k_cache=packed_index_k, + page_size=page_size, + ) + expected_indices, expected_lens = _reference_prefill_topk( + q, + weights, + index_k, + kv_workspace_slots, + row_starts, + row_ends, + topk=topk, + softmax_scale=softmax_scale, + ) + + _assert_topk_matches(workspace_indices, topk_lens, expected_indices, expected_lens) + + +def test_dsa_prefill_select_topk_keeps_late_values_above_threshold() -> None: + device = "cuda" + cols = 16384 + topk = 2048 + logits = torch.full((1, cols), -10.0, device=device, dtype=torch.float32) + equal_indices = torch.arange(0, 32, device=device, dtype=torch.int32) + greater_indices = torch.cat( + ( + torch.arange(4096, 4096 + topk - 2, device=device, dtype=torch.int32), + torch.tensor([cols - 3], device=device, dtype=torch.int32), + ) + ) + logits[0, equal_indices.long()] = 1.0 + logits[0, greater_indices.long()] = 2.0 + row_starts = torch.tensor([0], device=device, dtype=torch.int32) + row_ends = torch.tensor([cols], device=device, dtype=torch.int32) + out = torch.empty((1, topk), device=device, dtype=torch.int32) + lens_out = torch.empty((1,), device=device, dtype=torch.int32) + + num_warps = 8 + block_n = dsa_topk_gfx950.triton.next_power_of_2(cols) + dsa_topk_gfx950._dsa_prefill_select_topk_kernel[(1,)]( + logits, + row_starts, + row_ends, + out, + lens_out, + logits.stride(0), + out.stride(0), + topk=topk, + BLOCK_N=block_n, + LOAD_ELEMS=dsa_topk_gfx950._load_elems(block_n, num_warps), + TOPK_LOAD_ELEMS=dsa_topk_gfx950._load_elems(topk, num_warps), + num_warps=num_warps, + ) + torch.cuda.synchronize() + + selected = out[0, :topk] + selected_set = set(selected.cpu().tolist()) + assert selected_set.issuperset(set(greater_indices.cpu().tolist())) + assert len(selected_set.intersection(set(equal_indices.cpu().tolist()))) == 1 + torch.testing.assert_close(lens_out.cpu(), torch.tensor([topk], dtype=torch.int32)) + + +def test_dsa_decode_select_topk_keeps_late_values_above_threshold() -> None: + device = "cuda" + page_size = 64 + cols = 16384 + topk = 2048 + logits = torch.full((1, cols), -10.0, device=device, dtype=torch.float32) + equal_indices = torch.arange(0, 32, device=device, dtype=torch.int32) + greater_indices = torch.cat( + ( + torch.arange(4096, 4096 + topk - 2, device=device, dtype=torch.int32), + torch.tensor([cols - 3], device=device, dtype=torch.int32), + ) + ) + logits[0, equal_indices.long()] = 1.0 + logits[0, greater_indices.long()] = 2.0 + seq_lens = torch.tensor([cols], device=device, dtype=torch.int32) + block_table = torch.arange( + math.ceil(cols / page_size), device=device, dtype=torch.int32 + ).reshape(1, -1) + out = torch.empty((1, topk), device=device, dtype=torch.int32) + lens_out = torch.empty((1,), device=device, dtype=torch.int32) + + num_warps = 8 + block_n = dsa_topk_gfx950.triton.next_power_of_2(cols) + dsa_topk_gfx950._dsa_decode_select_topk_kernel[(1,)]( + logits, + block_table, + seq_lens, + out, + lens_out, + logits.stride(0), + block_table.stride(0), + out.stride(0), + block_table.shape[1], + page_size=page_size, + topk=topk, + q_len_per_req=1, + BLOCK_N=block_n, + LOAD_ELEMS=dsa_topk_gfx950._load_elems(block_n, num_warps), + TOPK_LOAD_ELEMS=dsa_topk_gfx950._load_elems(topk, num_warps), + num_warps=num_warps, + ) + torch.cuda.synchronize() + + selected = out[0, :topk] + selected_set = set(selected.cpu().tolist()) + assert selected_set.issuperset(set(greater_indices.cpu().tolist())) + assert len(selected_set.intersection(set(equal_indices.cpu().tolist()))) == 1 + torch.testing.assert_close(lens_out.cpu(), torch.tensor([topk], dtype=torch.int32)) + + +def test_dsa_decode_topk_gluon_long_row_uses_radix_path() -> None: + device = "cuda" + page_size = 64 + seq_len = 65536 + topk = 2048 + head_dim = 128 + q = torch.ones((1, 1, head_dim), device=device, dtype=torch.bfloat16) + weights = torch.ones((1, 1), device=device, dtype=torch.float32) + index_k = torch.zeros((seq_len, head_dim), device=device, dtype=torch.bfloat16) + index_k[:topk].fill_(1.0) + packed_index_k, _ = _pack_index_k_cache(index_k, page_size) + seq_lens = torch.tensor([seq_len], device=device, dtype=torch.int32) + block_table = torch.arange( + seq_len // page_size, device=device, dtype=torch.int32 + ).reshape(1, -1) + + topk_slots, topk_lens = gluon_dsa_decode_topk_fp8_gfx950( + q, + weights, + seq_lens, + block_table, + page_size=page_size, + topk=topk, + softmax_scale=head_dim**-0.5, + index_k_cache=packed_index_k, + ) + + expected = torch.arange(topk, device=device, dtype=torch.int32) + torch.testing.assert_close(topk_lens.cpu(), torch.tensor([topk], dtype=torch.int32)) + torch.testing.assert_close(torch.sort(topk_slots[0]).values.cpu(), expected.cpu()) + + +def test_dsa_prefill_topk_gluon_long_row_uses_radix_path() -> None: + device = "cuda" + page_size = 64 + seq_len = 65536 + topk = 2048 + head_dim = 128 + q = torch.ones((1, 1, head_dim), device=device, dtype=torch.bfloat16) + weights = torch.ones((1, 1), device=device, dtype=torch.float32) + index_k = torch.zeros((seq_len, head_dim), device=device, dtype=torch.bfloat16) + index_k[:topk].fill_(1.0) + packed_index_k, _ = _pack_index_k_cache(index_k, page_size) + kv_workspace_slots = torch.arange(seq_len, device=device, dtype=torch.int64) + row_starts = torch.tensor([0], device=device, dtype=torch.int32) + row_ends = torch.tensor([seq_len], device=device, dtype=torch.int32) + + workspace_indices, topk_lens = gluon_dsa_prefill_topk_fp8_gfx950( + q, + weights, + kv_workspace_slots, + row_starts, + row_ends, + topk=topk, + softmax_scale=head_dim**-0.5, + index_k_cache=packed_index_k, + page_size=page_size, + ) + + expected = torch.arange(topk, device=device, dtype=torch.int32) + torch.testing.assert_close(topk_lens.cpu(), torch.tensor([topk], dtype=torch.int32)) + torch.testing.assert_close( + torch.sort(workspace_indices[0]).values.cpu(), + expected.cpu(), + ) + + +def _pack_sparse_kv( + latent: torch.Tensor, + rope: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + kv_lora_rank = latent.shape[1] + qk_rope_head_dim = rope.shape[1] + scale = latent.float().abs().amax(dim=1, keepdim=True).clamp_min(1.0e-6) / 448.0 + latent_fp8 = (latent.float() / scale).clamp(-448.0, 448.0).to(torch.float8_e4m3fn) + row_bytes = kv_lora_rank + kv_lora_rank // 128 * 4 + qk_rope_head_dim * 2 + sparse = torch.empty( + (latent.shape[0], row_bytes), + dtype=torch.uint8, + device=latent.device, + ) + sparse[:, :kv_lora_rank].copy_(latent_fp8.view(torch.uint8)) + scale_start = kv_lora_rank + scale_end = scale_start + kv_lora_rank // 128 * 4 + sparse[:, scale_start:scale_end].view(torch.float32).copy_(scale) + sparse[:, scale_end:].view(torch.bfloat16).copy_(rope) + return sparse, latent_fp8.float() * scale + + +def _dsa_reference( + q: torch.Tensor, + latent: torch.Tensor, + rope: torch.Tensor, + topk_slots: torch.Tensor, + topk_lens: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + refs = [] + kv_lora_rank = latent.shape[1] + for token in range(q.shape[0]): + valid_slots = topk_slots[token, : int(topk_lens[token].item())].long() + valid_slots = valid_slots[valid_slots >= 0] + q_nope = q[token, :, :kv_lora_rank].float() + q_rope = q[token, :, kv_lora_rank:].float() + if valid_slots.numel() == 0: + refs.append(torch.zeros_like(q_nope)) + continue + k_nope = latent.index_select(0, valid_slots).float() + k_rope = rope.index_select(0, valid_slots).float() + scores = torch.einsum("hd,kd->hk", q_nope, k_nope) + scores += torch.einsum("hd,kd->hk", q_rope, k_rope) + probs = torch.softmax(scores * softmax_scale, dim=-1) + refs.append(torch.matmul(probs, k_nope)) + return torch.stack(refs, dim=0).to(torch.bfloat16) + + +def _dsa_visible_ranges(case: _DSACase, device: str) -> tuple[list[range], int]: + if case.mode == "decode": + assert case.visible_lens is not None + ranges = [range(0, int(visible_len)) for visible_len in case.visible_lens] + return ranges, _round_up_to_page(max(case.visible_lens), 64) + + assert case.prefix_lens is not None + assert case.extend_lens is not None + kv_workspace_slots, _, _, ranges = _make_prefill_workspace( + case.prefix_lens, + case.extend_lens, + device=device, + ) + return ranges, _round_up_to_page(int(kv_workspace_slots.numel()), 64) + + +def _make_selected_topk_slots( + case: _DSACase, + visible_ranges: Sequence[range], + *, + device: str, + generator: torch.Generator, +) -> tuple[torch.Tensor, torch.Tensor]: + assert case.topk_lens is not None + assert len(case.topk_lens) == len(visible_ranges) + topk_slots = torch.full( + (len(visible_ranges), case.topk), -1, device=device, dtype=torch.int32 + ) + lens: list[int] = [] + for token, visible_range in enumerate(visible_ranges): + visible_count = len(visible_range) + count = min(int(case.topk_lens[token]), visible_count, int(case.topk)) + lens.append(count) + if count == 0: + continue + candidates = torch.arange( + visible_range.start, + visible_range.stop, + device=device, + dtype=torch.int32, + ) + perm = torch.randperm(visible_count, device=device, generator=generator)[:count] + topk_slots[token, :count] = candidates.index_select(0, perm) + return topk_slots, torch.tensor(lens, device=device, dtype=torch.int32) + + +def _assert_slots_visible( + topk_slots: torch.Tensor, + topk_lens: torch.Tensor, + visible_ranges: Sequence[range], +) -> None: + for token, visible_range in enumerate(visible_ranges): + count = int(topk_lens[token].item()) + valid = topk_slots[token, :count] + if count: + assert (valid >= visible_range.start).all() + assert (valid < visible_range.stop).all() + assert (topk_slots[token, count:] == -1).all() + + +@pytest.mark.parametrize( + "mode,api", + [ + pytest.param("decode", gluon_dsa_decode_gfx950, id="decode"), + pytest.param("prefill", gluon_dsa_prefill_gfx950, id="prefill"), + ], +) +@pytest.mark.parametrize( + "q_dtype", + [ + pytest.param(torch.bfloat16, id="q_bf16"), + pytest.param(torch.float8_e4m3fn, id="q_fp8"), + ], +) +def test_dsa_with_sparse_kvcache(mode: str, api, q_dtype: torch.dtype) -> None: + device = "cuda" + tokens = 3 + num_heads = 2 + num_slots = 16 + topk = 512 + kv_lora_rank = 128 + qk_rope_head_dim = 64 + qk_nope_head_dim = 128 + softmax_scale = 1.0 / math.sqrt(qk_nope_head_dim + qk_rope_head_dim) + q_bf16 = torch.randn( + tokens, + num_heads, + kv_lora_rank + qk_rope_head_dim, + device=device, + dtype=torch.bfloat16, + ) + q = q_bf16.to(q_dtype) + latent = torch.randn(num_slots, kv_lora_rank, device=device, dtype=torch.bfloat16) + rope = torch.randn(num_slots, qk_rope_head_dim, device=device, dtype=torch.bfloat16) + sparse_kv, dequant_latent = _pack_sparse_kv(latent, rope) + topk_slots = torch.full((tokens, topk), -1, device=device, dtype=torch.int32) + topk_lens = torch.tensor([5, 7, 4], device=device, dtype=torch.int32) + for token in range(tokens): + count = int(topk_lens[token].item()) + topk_slots[token, :count] = torch.randperm(num_slots, device=device)[:count] + + out = api( + q=q, + kv_cache=None, + sparse_kv_cache=sparse_kv, + topk_slots=topk_slots, + topk_lens=topk_lens, + max_seqlen_k=num_slots, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + softmax_scale=softmax_scale, + page_size=64, + ) + + ref = _dsa_reference( + q, + dequant_latent, + rope, + topk_slots, + topk_lens, + softmax_scale, + ) + assert out.shape == (tokens, num_heads, kv_lora_rank) + assert out.dtype == torch.bfloat16 + torch.testing.assert_close(out.float(), ref.float(), rtol=8e-2, atol=8e-2) + + +@pytest.mark.parametrize( + "mode,api", + [ + pytest.param("decode", gluon_dsa_decode_gfx950, id="decode"), + pytest.param("prefill", gluon_dsa_prefill_gfx950, id="prefill"), + ], +) +@pytest.mark.parametrize( + "q_dtype", + [ + pytest.param(torch.bfloat16, id="q_bf16"), + pytest.param(torch.float8_e4m3fn, id="q_fp8"), + ], +) +def test_dsa_dense_kvcache(mode: str, api, q_dtype: torch.dtype) -> None: + device = "cuda" + tokens = 3 + num_heads = 2 + num_slots = 16 + topk = 512 + kv_lora_rank = 128 + qk_rope_head_dim = 64 + qk_nope_head_dim = 128 + softmax_scale = 1.0 / math.sqrt(qk_nope_head_dim + qk_rope_head_dim) + q_bf16 = torch.randn( + tokens, + num_heads, + kv_lora_rank + qk_rope_head_dim, + device=device, + dtype=torch.bfloat16, + ) + q = q_bf16.to(q_dtype) + latent = torch.randn(num_slots, kv_lora_rank, device=device, dtype=torch.bfloat16) + rope = torch.randn(num_slots, qk_rope_head_dim, device=device, dtype=torch.bfloat16) + kv_cache = torch.cat([latent, rope], dim=-1).to(q_dtype) + dequant_latent = kv_cache[:, :kv_lora_rank].float().to(torch.bfloat16) + dequant_rope = kv_cache[:, kv_lora_rank:].float().to(torch.bfloat16) + topk_slots = torch.full((tokens, topk), -1, device=device, dtype=torch.int32) + topk_lens = torch.tensor([5, 7, 4], device=device, dtype=torch.int32) + for token in range(tokens): + count = int(topk_lens[token].item()) + topk_slots[token, :count] = torch.randperm(num_slots, device=device)[:count] + + out = api( + q=q, + kv_cache=kv_cache, + sparse_kv_cache=None, + topk_slots=topk_slots, + topk_lens=topk_lens, + max_seqlen_k=num_slots, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + softmax_scale=softmax_scale, + page_size=64, + ) + + ref = _dsa_reference( + q, + dequant_latent, + dequant_rope, + topk_slots, + topk_lens, + softmax_scale, + ) + assert out.shape == (tokens, num_heads, kv_lora_rank) + assert out.dtype == torch.bfloat16 + torch.testing.assert_close(out.float(), ref.float(), rtol=8e-2, atol=8e-2) + + +def test_dsa_decode_sparse_kvcache_trims_large_topk_for_tiny_lens() -> None: + device = "cuda" + tokens = 2 + num_heads = 2 + num_slots = 8 + topk = 2048 + kv_lora_rank = 128 + qk_rope_head_dim = 64 + qk_nope_head_dim = 128 + softmax_scale = 1.0 / math.sqrt(qk_nope_head_dim + qk_rope_head_dim) + q = torch.randn( + tokens, + num_heads, + kv_lora_rank + qk_rope_head_dim, + device=device, + dtype=torch.bfloat16, + ).to(torch.float8_e4m3fn) + latent = torch.randn(num_slots, kv_lora_rank, device=device, dtype=torch.bfloat16) + rope = torch.randn(num_slots, qk_rope_head_dim, device=device, dtype=torch.bfloat16) + sparse_kv, dequant_latent = _pack_sparse_kv(latent, rope) + topk_slots = torch.full((tokens, topk), -1, device=device, dtype=torch.int32) + topk_lens = torch.tensor([1, 3], device=device, dtype=torch.int32) + topk_slots[0, 0] = 2 + topk_slots[1, :3] = torch.tensor([0, 5, 7], device=device, dtype=torch.int32) + + out = gluon_dsa_decode_gfx950( + q=q, + kv_cache=None, + sparse_kv_cache=sparse_kv, + topk_slots=topk_slots, + topk_lens=topk_lens, + max_seqlen_k=num_slots, + qk_nope_head_dim=qk_nope_head_dim, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + softmax_scale=softmax_scale, + page_size=64, + ) + + ref = _dsa_reference( + q, + dequant_latent, + rope, + topk_slots, + topk_lens, + softmax_scale, + ) + assert out.shape == (tokens, num_heads, kv_lora_rank) + assert out.dtype == torch.bfloat16 + torch.testing.assert_close(out.float(), ref.float(), rtol=8e-2, atol=8e-2) + + +@pytest.mark.parametrize( + "case", + _GLM52_DSA_CASES, + ids=lambda case: case.name, +) +def test_dsa_glm52_selected_attention_cases(case: _DSACase) -> None: + device = "cuda" + page_size = 64 + gen = _generator(device, case.seed) + visible_ranges, num_slots = _dsa_visible_ranges(case, device) + tokens = len(visible_ranges) + q = _randn_bf16( + (tokens, case.num_heads, case.kv_lora_rank + case.qk_rope_head_dim), + device=device, + generator=gen, + ) + latent = _randn_bf16((num_slots, case.kv_lora_rank), device=device, generator=gen) + rope = _randn_bf16((num_slots, case.qk_rope_head_dim), device=device, generator=gen) + topk_slots, topk_lens = _make_selected_topk_slots( + case, visible_ranges, device=device, generator=gen + ) + _assert_slots_visible(topk_slots, topk_lens, visible_ranges) + + kv_cache = None + sparse_kv_cache = None + if case.kv_layout == "dense": + kv_cache = torch.cat([latent, rope], dim=-1).contiguous() + reference_latent = kv_cache[:, : case.kv_lora_rank] + reference_rope = kv_cache[:, case.kv_lora_rank :] + elif case.kv_layout == "sparse": + sparse_kv_cache, reference_latent = _pack_sparse_kv(latent, rope) + reference_rope = rope + else: + raise AssertionError(f"unknown DSA KV layout {case.kv_layout!r}") + + softmax_scale = 1.0 / math.sqrt(case.qk_nope_head_dim + case.qk_rope_head_dim) + common_kwargs = { + "q": q, + "kv_cache": kv_cache, + "sparse_kv_cache": sparse_kv_cache, + "topk_slots": topk_slots, + "topk_lens": topk_lens, + "max_seqlen_k": max(len(visible_range) for visible_range in visible_ranges), + "qk_nope_head_dim": case.qk_nope_head_dim, + "kv_lora_rank": case.kv_lora_rank, + "qk_rope_head_dim": case.qk_rope_head_dim, + "softmax_scale": softmax_scale, + "page_size": page_size, + } + if case.mode == "decode": + out = gluon_dsa_decode_gfx950(q_len_per_req=case.q_len_per_req, **common_kwargs) + else: + out = gluon_dsa_prefill_gfx950(**common_kwargs) + + ref = _dsa_reference( + q, + reference_latent, + reference_rope, + topk_slots, + topk_lens, + softmax_scale, + ) + assert out.shape == (tokens, case.num_heads, case.kv_lora_rank) + assert out.dtype == torch.bfloat16 + torch.testing.assert_close(out.float(), ref.float(), rtol=8e-2, atol=8e-2) diff --git a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py index b95e81947..b286116b4 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py @@ -29,9 +29,28 @@ current_platform, ) from tokenspeed_kernel.registry import Priority, register_kernel -from tokenspeed_kernel.signature import format_signatures +from tokenspeed_kernel.signature import ( + dense_tensor_format, + format_signature, + format_signatures, +) if current_platform().is_amd: + _DSA_FULL_TOPK_WIDTHS = frozenset({512, 1024, 2048}) + _DSA_PREFILL_TOPK_WIDTHS = _DSA_FULL_TOPK_WIDTHS + + from tokenspeed_kernel_amd.ops.attention.gluon.dsa_gfx950 import ( + gluon_dsa_decode_gfx950 as _dsa_decode_impl, + ) + from tokenspeed_kernel_amd.ops.attention.gluon.dsa_gfx950 import ( + gluon_dsa_prefill_gfx950 as _dsa_prefill_impl, + ) + from tokenspeed_kernel_amd.ops.attention.gluon.dsa_topk_gfx950 import ( + gluon_dsa_decode_topk_fp8_gfx950 as _dsa_decode_topk_impl, + ) + from tokenspeed_kernel_amd.ops.attention.gluon.dsa_topk_gfx950 import ( + gluon_dsa_prefill_topk_fp8_gfx950 as _dsa_prefill_topk_impl, + ) from tokenspeed_kernel_amd.ops.attention.gluon.mha_decode_gfx950 import ( gluon_mha_decode_gfx950 as _decode_impl, ) @@ -140,3 +159,128 @@ def gluon_mha_prefill_gfx950(*args, **kwargs): ) def gluon_mha_extend_gfx950(*args, **kwargs): return _extend_impl(*args, **kwargs) + + @register_kernel( + "attention", + "dsa_decode_topk", + name="gluon_dsa_decode_topk_fp8_gfx950", + solution="gluon", + capability=CapabilityRequirement( + min_arch_version=ArchVersion(9, 5), + max_arch_version=ArchVersion(9, 5), + vendors=frozenset({"amd"}), + ), + signatures=frozenset( + { + format_signature( + q=dense_tensor_format(torch.bfloat16), + weights=dense_tensor_format(torch.float32), + ) + } + ), + priority=Priority.SPECIALIZED, + traits={ + "head_dim": frozenset({128}), + "topk": frozenset({512, 1024, 2048}), + "page_size": frozenset({64}), + "q_len_per_req": frozenset({1, 2, 3, 4, 5, 6}), + "index_k_format": frozenset({"fp8_scaled"}), + }, + ) + def gluon_dsa_decode_topk_fp8_gfx950(*args, **kwargs): + return _dsa_decode_topk_impl(*args, **kwargs) + + @register_kernel( + "attention", + "dsa_prefill_topk", + name="gluon_dsa_prefill_topk_fp8_gfx950", + solution="gluon", + capability=CapabilityRequirement( + min_arch_version=ArchVersion(9, 5), + max_arch_version=ArchVersion(9, 5), + vendors=frozenset({"amd"}), + ), + signatures=frozenset( + { + format_signature( + q=dense_tensor_format(torch.bfloat16), + weights=dense_tensor_format(torch.float32), + ) + } + ), + priority=Priority.SPECIALIZED, + traits={ + "head_dim": frozenset({128}), + "topk": frozenset({512, 1024, 2048}), + "index_k_format": frozenset({"fp8_scaled"}), + }, + ) + def gluon_dsa_prefill_topk_fp8_gfx950(*args, **kwargs): + return _dsa_prefill_topk_impl(*args, **kwargs) + + @register_kernel( + "attention", + "dsa_decode", + name="gluon_dsa_decode_gfx950", + solution="gluon", + capability=CapabilityRequirement( + min_arch_version=ArchVersion(9, 5), + max_arch_version=ArchVersion(9, 5), + vendors=frozenset({"amd"}), + ), + signatures=frozenset( + { + format_signature(q=dense_tensor_format(torch.bfloat16)), + format_signature(q=dense_tensor_format(torch.float8_e4m3fn)), + } + ), + priority=Priority.SPECIALIZED, + traits={ + "page_size": frozenset({64}), + "q_len_per_req": frozenset({1, 2, 3, 4, 5, 6}), + "qk_nope_head_dim": frozenset({128, 192}), + "kv_lora_rank": frozenset({128, 512}), + "qk_rope_head_dim": frozenset({64}), + "topk": _DSA_FULL_TOPK_WIDTHS, + "kv_cache_available": frozenset({False, True}), + "sparse_kv_cache_available": frozenset({False, True}), + "topk_layout": frozenset({"global_slots"}), + "support_logit_cap": frozenset({False}), + "return_lse": frozenset({False}), + }, + ) + def gluon_dsa_decode_gfx950(*args, **kwargs): + return _dsa_decode_impl(*args, **kwargs) + + @register_kernel( + "attention", + "dsa_prefill", + name="gluon_dsa_prefill_gfx950", + solution="gluon", + capability=CapabilityRequirement( + min_arch_version=ArchVersion(9, 5), + max_arch_version=ArchVersion(9, 5), + vendors=frozenset({"amd"}), + ), + signatures=frozenset( + { + format_signature(q=dense_tensor_format(torch.bfloat16)), + } + ), + priority=Priority.SPECIALIZED, + traits={ + "page_size": frozenset({64}), + "q_len_per_req": frozenset({1}), + "qk_nope_head_dim": frozenset({128, 192}), + "kv_lora_rank": frozenset({128, 512}), + "qk_rope_head_dim": frozenset({64}), + "topk": _DSA_PREFILL_TOPK_WIDTHS, + "kv_cache_available": frozenset({False, True}), + "sparse_kv_cache_available": frozenset({False, True}), + "topk_layout": frozenset({"global_slots"}), + "support_logit_cap": frozenset({False}), + "return_lse": frozenset({False}), + }, + ) + def gluon_dsa_prefill_gfx950(*args, **kwargs): + return _dsa_prefill_impl(*args, **kwargs) diff --git a/tokenspeed-kernel/test/test_kernel_api_selection.py b/tokenspeed-kernel/test/test_kernel_api_selection.py index 111c30b52..354c23b73 100644 --- a/tokenspeed-kernel/test/test_kernel_api_selection.py +++ b/tokenspeed-kernel/test/test_kernel_api_selection.py @@ -508,7 +508,6 @@ def _attention_dsa_decode() -> object: qk_rope_head_dim=64, softmax_scale=1.0, page_size=64, - solution="triton", ) @@ -529,7 +528,26 @@ def _attention_dsa_prefill() -> object: qk_rope_head_dim=64, softmax_scale=1.0, page_size=64, - solution="triton", + ) + + +def _attention_dsa_prefill_fp8_dense() -> object: + q = torch.empty((2, 8, 576), dtype=torch.float8_e4m3fn) + kv_cache = torch.empty((64, 576), dtype=torch.float8_e4m3fn) + topk_slots = torch.empty((2, 1024), dtype=torch.int32) + topk_lens = torch.empty((2,), dtype=torch.int32) + return tokenspeed_kernel.dsa_prefill( + q=q, + kv_cache=kv_cache, + sparse_kv_cache=None, + topk_slots=topk_slots, + topk_lens=topk_lens, + max_seqlen_k=64, + qk_nope_head_dim=192, + kv_lora_rank=512, + qk_rope_head_dim=64, + softmax_scale=1.0, + page_size=64, ) @@ -1295,7 +1313,7 @@ def _case( "cdna4", "attention", "dsa_decode", - "triton_dsa_decode", + "gluon_dsa_decode_gfx950", _attention_dsa_decode, ), _case( @@ -1303,15 +1321,23 @@ def _case( "cdna4", "attention", "dsa_prefill", - "triton_dsa_prefill", + "gluon_dsa_prefill_gfx950", _attention_dsa_prefill, ), + _case( + _is_cdna4, + "cdna4", + "attention", + "dsa_prefill", + "triton_dsa_prefill", + _attention_dsa_prefill_fp8_dense, + ), _case( _is_cdna4, "cdna4", "attention", "dsa_decode_topk", - "triton_dsa_decode_topk_fp8", + "gluon_dsa_decode_topk_fp8_gfx950", _attention_dsa_decode_topk, ), _case( @@ -1319,7 +1345,7 @@ def _case( "cdna4", "attention", "dsa_prefill_topk", - "triton_dsa_prefill_topk_fp8", + "gluon_dsa_prefill_topk_fp8_gfx950", _attention_dsa_prefill_topk, ), _case(