From 5b80ee416162e2cbee1714fcfe542f4e92f90d0c Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Wed, 8 Jul 2026 16:51:29 +0000 Subject: [PATCH 01/14] Add GLM-style DSA correctness tests Signed-off-by: Max Dawkins --- .../test/ops/test_attention_dsa.py | 674 ++++++++++++++++++ 1 file changed, 674 insertions(+) diff --git a/tokenspeed-kernel/test/ops/test_attention_dsa.py b/tokenspeed-kernel/test/ops/test_attention_dsa.py index 50e19c6ef..550978630 100644 --- a/tokenspeed-kernel/test/ops/test_attention_dsa.py +++ b/tokenspeed-kernel/test/ops/test_attention_dsa.py @@ -21,6 +21,8 @@ from __future__ import annotations import math +from collections.abc import Sequence +from dataclasses import dataclass import pytest import torch @@ -38,6 +40,246 @@ torch.manual_seed(42) +@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, + ), +) + + def _pack_index_k_cache( index_k: torch.Tensor, page_size: int, @@ -74,6 +316,177 @@ def _pack_index_k_cache( 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, +) -> tuple[torch.Tensor, torch.Tensor]: + out = torch.full((q.shape[0], topk), -1, device=q.device, dtype=torch.int32) + lens = torch.minimum(seq_lens.to(torch.int32), torch.full_like(seq_lens, int(topk))) + for token in range(q.shape[0]): + seq_len = int(seq_lens[token].item()) + count = int(lens[token].item()) + if count == 0: + continue + offsets = torch.arange(seq_len, device=q.device, dtype=torch.long) + pages = block_table[token].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 test_dsa_decode_topk_fp8(device: str, require) -> None: require("attention", "dsa_decode_topk", "triton", torch.bfloat16, "q") @@ -234,6 +647,127 @@ def test_dsa_prefill_topk_fp8(device: str, require) -> None: assert (workspace_indices[0, int(expected_lens[0].item()) :] == -1).all() +@pytest.mark.parametrize("solution", ["triton"]) +@pytest.mark.parametrize( + "case", + _GLM52_TOPK_DECODE_CASES, + ids=lambda case: case.name, +) +def test_dsa_decode_topk_fp8_glm52_cases( + device: str, + require, + solution: str, + case: _TopKDecodeCase, +) -> None: + require("attention", "dsa_decode_topk", solution, torch.bfloat16, "q") + + 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) + q = _randn_bf16( + (len(case.seq_lens), case.index_heads, head_dim), + device=device, + generator=gen, + ) + weights = _normal_weights( + (len(case.seq_lens), 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 = dsa_decode_topk( + q, + weights, + seq_lens, + block_table, + page_size=page_size, + topk=case.topk, + softmax_scale=softmax_scale, + q_len_per_req=case.q_len_per_req, + index_k_cache=packed_index_k, + solution=solution, + ) + 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, + ) + + _assert_topk_matches(topk_slots, topk_lens, expected_slots, expected_lens) + + +@pytest.mark.parametrize("solution", ["triton"]) +@pytest.mark.parametrize( + "case", + _GLM52_TOPK_PREFILL_CASES, + ids=lambda case: case.name, +) +def test_dsa_prefill_topk_fp8_glm52_cases( + device: str, + require, + solution: str, + case: _TopKPrefillCase, +) -> None: + require("attention", "dsa_prefill_topk", solution, torch.bfloat16, "q") + + 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 = dsa_prefill_topk( + 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, + solution=solution, + ) + 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_plan_triton(device: str) -> None: # The triton decode kernel derives its own causal bounds and ignores the # plan, so triton_dsa_plan is a no-op returning an opaque, non-None @@ -317,8 +851,12 @@ def _dsa_reference( 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) @@ -328,6 +866,66 @@ def _dsa_reference( 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_name", [ @@ -470,3 +1068,79 @@ def test_dsa_decode_dense_kvcache(device: str, q_dtype: torch.dtype, require) -> 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("solution", ["triton"]) +@pytest.mark.parametrize( + "case", + _GLM52_DSA_CASES, + ids=lambda case: case.name, +) +def test_dsa_glm52_selected_attention_cases( + device: str, + require, + solution: str, + case: _DSACase, +) -> None: + api_name = "dsa_decode" if case.mode == "decode" else "dsa_prefill" + require("attention", api_name, solution, torch.bfloat16, "q") + + 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, + "solution": solution, + } + if case.mode == "decode": + out = dsa_decode(q_len_per_req=case.q_len_per_req, **common_kwargs) + else: + out = dsa_prefill(**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) From f774151c27316a99f7663ba97b7a7d655a4f2563 Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Wed, 8 Jul 2026 16:57:51 +0000 Subject: [PATCH 02/14] Add Gluon DSA top-k path Signed-off-by: Max Dawkins --- .../ops/attention/gluon/__init__.py | 4 + .../ops/attention/gluon/dsa_topk_gfx950.py | 482 ++++++++++++++++++ .../ops/attention/gluon/__init__.py | 70 ++- .../test/ops/test_attention_dsa.py | 4 +- 4 files changed, 557 insertions(+), 3 deletions(-) create mode 100644 tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_topk_gfx950.py 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..c841ff5f8 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 @@ -25,6 +25,10 @@ from tokenspeed_kernel_amd.ops.attention.gluon.mha_decode_gfx950 import ( # noqa: F401 gluon_mha_decode_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_extend_gfx950 import ( # noqa: F401 gluon_mha_extend_gfx950, ) 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..b18b43ce9 --- /dev/null +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_topk_gfx950.py @@ -0,0 +1,482 @@ +# 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. + +"""Correctness-first DSA top-k Gluon kernels for AMD GFX950.""" + +from __future__ import annotations + +import torch +from tokenspeed_kernel_amd._triton import gl, gluon, triton + +__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 _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, + BLOCK_N: gl.constexpr, +): + token = gl.program_id(0) + block_id = gl.program_id(1) + offsets = block_id * BLOCK_N + gl.arange( + 0, BLOCK_N, layout=_vector_layout(BLOCK_N, gl.num_warps(), 1) + ) + seq_len = gl.load(seq_lens + token).to(gl.int32) + valid = (offsets < seq_len) & (offsets < max_seq_len) + block_idx = offsets // page_size + block_offset = offsets - block_idx * page_size + page = gl.load( + block_table + token * block_table_stride + block_idx, + 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=_vector_layout(BLOCK_N, gl.num_warps(), 1), + ) + + 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=_vector_layout(BLOCK_N, gl.num_warps(), 1), + ) + for dim in gl.static_range(0, head_dim): + q_val = gl.load(q + (token * num_heads + head) * head_dim + dim).to( + gl.float32 + ) + k_val = gl.load( + index_k_fp8 + fp8_base + dim, + mask=valid, + other=0.0, + ).to(gl.float32) + k_scale = gl.load( + index_k_scale + scale_base + dim // 128, + mask=valid, + other=0.0, + ).to(gl.float32) + head_score += k_val * k_scale * q_val + 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, +): + token = gl.program_id(0) + block_id = gl.program_id(1) + offsets = block_id * BLOCK_N + gl.arange( + 0, BLOCK_N, layout=_vector_layout(BLOCK_N, gl.num_warps(), 1) + ) + 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.load(kv_workspace_slots + offsets, 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=_vector_layout(BLOCK_N, gl.num_warps(), 1), + ) + + 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=_vector_layout(BLOCK_N, gl.num_warps(), 1), + ) + for dim in gl.static_range(0, head_dim): + q_val = gl.load(q + (token * num_heads + head) * head_dim + dim).to( + gl.float32 + ) + k_val = gl.load( + index_k_fp8 + fp8_base + dim, + mask=valid, + other=0.0, + ).to(gl.float32) + k_scale = gl.load( + index_k_scale + scale_base + dim // 128, + mask=valid, + other=0.0, + ).to(gl.float32) + head_score += k_val * k_scale * q_val + 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 + + +def _topk_with_padding(logits: torch.Tensor, topk: int) -> torch.Tensor: + topk = int(topk) + if topk <= 0: + raise ValueError(f"topk must be positive, got {topk}") + out = torch.full( + (logits.shape[0], topk), -1, device=logits.device, dtype=torch.int32 + ) + count = min(topk, logits.shape[1]) + if count == 0: + return out + out[:, :count].copy_(torch.topk(logits, count, dim=1).indices.to(torch.int32)) + return out + + +def _local_topk_to_global_slots( + local_topk_offsets: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + *, + page_size: int, + out: torch.Tensor | None = None, + lens_out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + tokens, topk = local_topk_offsets.shape + if out is None: + out = torch.empty_like(local_topk_offsets) + if lens_out is None: + lens_out = torch.empty( + (tokens,), dtype=torch.int32, device=local_topk_offsets.device + ) + seq_lens = seq_lens.to(device=local_topk_offsets.device, dtype=torch.int32) + block_table = block_table.to(device=local_topk_offsets.device, dtype=torch.int32) + lens_out.copy_(torch.minimum(seq_lens, torch.full_like(seq_lens, int(topk)))) + + local = local_topk_offsets.to(torch.int64) + valid = (local >= 0) & (local < seq_lens[:, None].to(torch.int64)) + safe_local = torch.where(valid, local, torch.zeros_like(local)) + block_idx = torch.div(safe_local, int(page_size), rounding_mode="floor") + block_idx = block_idx.clamp(max=block_table.shape[1] - 1) + pages = torch.gather(block_table.to(torch.int64), 1, block_idx) + slots = pages * int(page_size) + safe_local.remainder(int(page_size)) + out.copy_(torch.where(valid, slots.to(torch.int32), torch.full_like(out, -1))) + return out, lens_out + + +def _workspace_topk_filter( + workspace_indices: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + *, + out: torch.Tensor, +) -> torch.Tensor: + valid = (workspace_indices >= row_starts[:, None]) & ( + workspace_indices < row_ends[:, None] + ) + out.copy_(torch.where(valid, workspace_indices, torch.full_like(out, -1))) + return 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, + plan: object | None = None, + out: torch.Tensor | None = None, + lens_out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + del plan, 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 or seq_lens.numel() != q.shape[0]: + raise ValueError( + "seq_lens must be [tokens], got " + f"{tuple(seq_lens.shape)} for q={tuple(q.shape)}" + ) + if block_table.dim() != 2 or block_table.shape[0] < q.shape[0]: + raise ValueError( + "block_table must have at least one row per token, 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 + if not q.is_cuda: + raise RuntimeError("DSA Gluon FP8 decode top-k requires CUDA tensors") + + q = q.contiguous() + index_k_cache = index_k_cache.contiguous() + weights = weights.contiguous() + seq_lens = seq_lens.to(device=q.device, dtype=torch.int32).contiguous() + block_table = block_table.to(device=q.device, dtype=torch.int32).contiguous() + max_seq_len = int(block_table.shape[1]) * int(page_size) + logits = torch.empty( + (q.shape[0], max_seq_len), dtype=torch.float32, device=q.device + ) + block_n = 64 + _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), + BLOCK_N=block_n, + num_warps=4, + ) + local_topk_offsets = _topk_with_padding(logits, int(topk)) + return _local_topk_to_global_slots( + local_topk_offsets, + block_table, + seq_lens, + page_size=int(page_size), + out=out, + lens_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 + 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], int(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) + out.fill_(-1) + lens_out.zero_() + if q.shape[0] == 0: + return out, lens_out + if not q.is_cuda: + raise RuntimeError("DSA Gluon FP8 prefill top-k requires CUDA tensors") + + q = q.contiguous() + index_k_cache = index_k_cache.contiguous() + weights = weights.contiguous() + kv_workspace_slots = kv_workspace_slots.to( + device=q.device, dtype=torch.int64 + ).contiguous() + row_starts = row_starts.to(device=q.device, dtype=torch.int32).contiguous() + row_ends = row_ends.to(device=q.device, dtype=torch.int32).contiguous() + seq_len_sum = int(kv_workspace_slots.numel()) + candidate_lens = (row_ends - row_starts).clamp_min(0) + lens_out.copy_( + torch.minimum(candidate_lens, torch.full_like(candidate_lens, int(topk))) + ) + if seq_len_sum == 0: + 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 = 64 + 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, + num_warps=4, + ) + workspace_indices = _topk_with_padding(logits, int(topk)) + _workspace_topk_filter( + workspace_indices, + row_starts[start:end], + row_ends[start:end], + out=out[start:end], + ) + return out, lens_out 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..de41a46d0 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py @@ -29,12 +29,22 @@ 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: from tokenspeed_kernel_amd.ops.attention.gluon.mha_decode_gfx950 import ( gluon_mha_decode_gfx950 as _decode_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_extend_gfx950 import ( gluon_mha_extend_gfx950 as _extend_impl, ) @@ -140,3 +150,61 @@ 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) diff --git a/tokenspeed-kernel/test/ops/test_attention_dsa.py b/tokenspeed-kernel/test/ops/test_attention_dsa.py index 550978630..4b85c08d4 100644 --- a/tokenspeed-kernel/test/ops/test_attention_dsa.py +++ b/tokenspeed-kernel/test/ops/test_attention_dsa.py @@ -647,7 +647,7 @@ def test_dsa_prefill_topk_fp8(device: str, require) -> None: assert (workspace_indices[0, int(expected_lens[0].item()) :] == -1).all() -@pytest.mark.parametrize("solution", ["triton"]) +@pytest.mark.parametrize("solution", ["triton", "gluon"]) @pytest.mark.parametrize( "case", _GLM52_TOPK_DECODE_CASES, @@ -706,7 +706,7 @@ def test_dsa_decode_topk_fp8_glm52_cases( _assert_topk_matches(topk_slots, topk_lens, expected_slots, expected_lens) -@pytest.mark.parametrize("solution", ["triton"]) +@pytest.mark.parametrize("solution", ["triton", "gluon"]) @pytest.mark.parametrize( "case", _GLM52_TOPK_PREFILL_CASES, From f38070cfc4c73adbf5575834f98f5c47327a8e40 Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Wed, 8 Jul 2026 17:13:10 +0000 Subject: [PATCH 03/14] Add Gluon DSA selected attention path Signed-off-by: Max Dawkins --- .../ops/attention/gluon/__init__.py | 4 + .../ops/attention/gluon/dsa_gfx950.py | 586 ++++++++++++++++++ .../ops/attention/gluon/__init__.py | 72 +++ .../test/ops/test_attention_dsa.py | 2 +- 4 files changed, 663 insertions(+), 1 deletion(-) create mode 100644 tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_gfx950.py 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 c841ff5f8..b95e849a9 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 @@ -25,6 +25,10 @@ from tokenspeed_kernel_amd.ops.attention.gluon.mha_decode_gfx950 import ( # noqa: F401 gluon_mha_decode_gfx950, ) +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, 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..eea35ceba --- /dev/null +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_gfx950.py @@ -0,0 +1,586 @@ +# 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. + +"""Correctness-first selected-slot DSA Gluon kernels for AMD GFX950.""" + +from __future__ import annotations + +import torch +from tokenspeed_kernel_amd._triton import gl, gluon, triton + +__all__ = [ + "gluon_dsa_decode_gfx950", + "gluon_dsa_prefill_gfx950", +] + + +@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_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 != torch.bfloat16: + raise TypeError(f"Gluon DSA supports BF16 q for this milestone, got {q.dtype}") + if page_size != 64: + raise ValueError(f"Gluon DSA supports page_size=64, got {page_size}") + if qk_nope_head_dim != 192: + raise ValueError( + f"Gluon DSA supports qk_nope_head_dim=192, got {qk_nope_head_dim}" + ) + if kv_lora_rank != 512: + raise ValueError(f"Gluon DSA supports kv_lora_rank=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}" + ) + 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 _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: + kv_dim = int(kv_lora_rank) + int(qk_rope_head_dim) + out = torch.empty( + (q.shape[0], q.shape[1], kv_lora_rank), dtype=q.dtype, 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=q.dtype, 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, +) -> 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() + softmax_scale = float(softmax_scale) * float(k_scale) + 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: + result = _run_dense_kv( + q, + _flatten_dense_kv_cache(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, + ) + 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 max_seqlen_k, 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, + ) + + +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 max_seqlen_k, 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, + ) 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 de41a46d0..55c9efe7e 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py @@ -39,6 +39,12 @@ from tokenspeed_kernel_amd.ops.attention.gluon.mha_decode_gfx950 import ( gluon_mha_decode_gfx950 as _decode_impl, ) + 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, ) @@ -208,3 +214,69 @@ def gluon_dsa_decode_topk_fp8_gfx950(*args, **kwargs): ) 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)), + } + ), + priority=Priority.SPECIALIZED, + traits={ + "page_size": frozenset({64}), + "q_len_per_req": frozenset({1, 2, 3, 4, 5, 6}), + "qk_nope_head_dim": frozenset({192}), + "kv_lora_rank": frozenset({512}), + "qk_rope_head_dim": frozenset({64}), + "topk": frozenset({512, 1024, 2048}), + "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({192}), + "kv_lora_rank": frozenset({512}), + "qk_rope_head_dim": frozenset({64}), + "topk": frozenset({512, 1024, 2048}), + "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/ops/test_attention_dsa.py b/tokenspeed-kernel/test/ops/test_attention_dsa.py index 4b85c08d4..270a14e6c 100644 --- a/tokenspeed-kernel/test/ops/test_attention_dsa.py +++ b/tokenspeed-kernel/test/ops/test_attention_dsa.py @@ -1070,12 +1070,12 @@ def test_dsa_decode_dense_kvcache(device: str, q_dtype: torch.dtype, require) -> torch.testing.assert_close(out.float(), ref.float(), rtol=8e-2, atol=8e-2) -@pytest.mark.parametrize("solution", ["triton"]) @pytest.mark.parametrize( "case", _GLM52_DSA_CASES, ids=lambda case: case.name, ) +@pytest.mark.parametrize("solution", ["triton", "gluon"]) def test_dsa_glm52_selected_attention_cases( device: str, require, From 4ada190e821410b2960b37607a6f5c3306243ea0 Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Wed, 8 Jul 2026 17:54:51 +0000 Subject: [PATCH 04/14] Port Gluon DSA dense path to tiled MFMA Signed-off-by: Max Dawkins --- .../ops/attention/gluon/dsa_gfx950.py | 453 +++++++++++++++++- 1 file changed, 451 insertions(+), 2 deletions(-) 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 index eea35ceba..4d87a83ae 100644 --- 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 @@ -18,12 +18,12 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -"""Correctness-first selected-slot DSA Gluon kernels for AMD GFX950.""" +"""Selected-slot DSA Gluon kernels for AMD GFX950.""" from __future__ import annotations import torch -from tokenspeed_kernel_amd._triton import gl, gluon, triton +from tokenspeed_kernel_amd._triton import gl, gluon, tl, triton __all__ = [ "gluon_dsa_decode_gfx950", @@ -124,6 +124,415 @@ def _packed_score( 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, @@ -394,6 +803,44 @@ def _run_dense_kv( 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=q.dtype, 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( @@ -483,6 +930,8 @@ def _run_dsa( topk_slots = topk_slots.contiguous() topk_lens = topk_lens.contiguous() 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, From 8248e61ac827138fd7521757c62a403c2c74e3e7 Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Wed, 8 Jul 2026 18:13:01 +0000 Subject: [PATCH 05/14] Implement Gluon DSA top-k selector Signed-off-by: Max Dawkins --- .../ops/attention/gluon/dsa_topk_gfx950.py | 301 ++++++++++++++---- 1 file changed, 231 insertions(+), 70 deletions(-) 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 index b18b43ce9..bf5d94749 100644 --- 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 @@ -18,7 +18,7 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -"""Correctness-first DSA top-k Gluon kernels for AMD GFX950.""" +"""DSA top-k Gluon kernels for AMD GFX950.""" from __future__ import annotations @@ -40,6 +40,185 @@ def _vector_layout( 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 = gl.full((), 0, dtype=gl.uint32) + remaining = gl.full((), topk, dtype=gl.int32) + + 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 = gl.full((), 0, dtype=gl.int32) + selected = gl.full((), 0, dtype=gl.uint32) + selected_remaining = remaining + found = gl.full((), 0, dtype=gl.int32) + + 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, + 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) + seq_len = gl.load(seq_lens + row).to(gl.int32) + 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 + row * 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) + selected = valid & (keys >= threshold) + selected_i32 = selected.to(gl.int32) + selected_pos = gl.associative_scan(selected_i32, 0, _topk_add) - 1 + write = selected & (selected_pos < topk) + local = offsets.to(gl.int32) + block_idx = local // page_size + block_offset = local - block_idx * page_size + page = gl.load( + block_table + row * block_table_stride + block_idx, + mask=write & (block_idx < block_table_cols), + other=0, + ).to(gl.int32) + slots = page * page_size + block_offset + gl.store(out + row * out_stride + selected_pos, slots, mask=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) + selected = valid & (keys >= threshold) + selected_i32 = selected.to(gl.int32) + selected_pos = gl.associative_scan(selected_i32, 0, _topk_add) - 1 + write = selected & (selected_pos < topk) + gl.store(out + row * out_stride + selected_pos, offsets.to(gl.int32), mask=write) + + @gluon.jit def _dsa_decode_logits_fp8_kernel( q, @@ -236,63 +415,21 @@ def _check_packed_fp8_inputs( return row_bytes -def _topk_with_padding(logits: torch.Tensor, topk: int) -> torch.Tensor: - topk = int(topk) - if topk <= 0: - raise ValueError(f"topk must be positive, got {topk}") - out = torch.full( - (logits.shape[0], topk), -1, device=logits.device, dtype=torch.int32 - ) - count = min(topk, logits.shape[1]) - if count == 0: - return out - out[:, :count].copy_(torch.topk(logits, count, dim=1).indices.to(torch.int32)) - return out +def _next_power_of_2(value: int) -> int: + if value <= 1: + return 1 + return 1 << (int(value) - 1).bit_length() -def _local_topk_to_global_slots( - local_topk_offsets: torch.Tensor, - block_table: torch.Tensor, - seq_lens: torch.Tensor, - *, - page_size: int, - out: torch.Tensor | None = None, - lens_out: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - tokens, topk = local_topk_offsets.shape - if out is None: - out = torch.empty_like(local_topk_offsets) - if lens_out is None: - lens_out = torch.empty( - (tokens,), dtype=torch.int32, device=local_topk_offsets.device - ) - seq_lens = seq_lens.to(device=local_topk_offsets.device, dtype=torch.int32) - block_table = block_table.to(device=local_topk_offsets.device, dtype=torch.int32) - lens_out.copy_(torch.minimum(seq_lens, torch.full_like(seq_lens, int(topk)))) - - local = local_topk_offsets.to(torch.int64) - valid = (local >= 0) & (local < seq_lens[:, None].to(torch.int64)) - safe_local = torch.where(valid, local, torch.zeros_like(local)) - block_idx = torch.div(safe_local, int(page_size), rounding_mode="floor") - block_idx = block_idx.clamp(max=block_table.shape[1] - 1) - pages = torch.gather(block_table.to(torch.int64), 1, block_idx) - slots = pages * int(page_size) + safe_local.remainder(int(page_size)) - out.copy_(torch.where(valid, slots.to(torch.int32), torch.full_like(out, -1))) - return out, lens_out +def _load_elems(block: int, num_warps: int) -> int: + return max(1, triton.cdiv(int(block), 64 * int(num_warps))) -def _workspace_topk_filter( - workspace_indices: torch.Tensor, - row_starts: torch.Tensor, - row_ends: torch.Tensor, - *, - out: torch.Tensor, -) -> torch.Tensor: - valid = (workspace_indices >= row_starts[:, None]) & ( - workspace_indices < row_ends[:, None] - ) - out.copy_(torch.where(valid, workspace_indices, torch.full_like(out, -1))) - return out +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 gluon_dsa_decode_topk_fp8_gfx950( @@ -311,6 +448,8 @@ def gluon_dsa_decode_topk_fp8_gfx950( lens_out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: del plan, q_len_per_req + topk = int(topk) + _validate_topk(topk) 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)) @@ -345,6 +484,10 @@ def gluon_dsa_decode_topk_fp8_gfx950( seq_lens = seq_lens.to(device=q.device, dtype=torch.int32).contiguous() block_table = block_table.to(device=q.device, dtype=torch.int32).contiguous() 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 ) @@ -369,15 +512,26 @@ def gluon_dsa_decode_topk_fp8_gfx950( BLOCK_N=block_n, num_warps=4, ) - local_topk_offsets = _topk_with_padding(logits, int(topk)) - return _local_topk_to_global_slots( - local_topk_offsets, + select_warps = 8 + select_block = _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), - out=out, - lens_out=lens_out, + 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 def gluon_dsa_prefill_topk_fp8_gfx950( @@ -398,6 +552,8 @@ def gluon_dsa_prefill_topk_fp8_gfx950( 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" @@ -414,11 +570,9 @@ def gluon_dsa_prefill_topk_fp8_gfx950( f"q={tuple(q.shape)}" ) if out is None: - out = torch.empty((q.shape[0], int(topk)), dtype=torch.int32, device=q.device) + 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) - out.fill_(-1) - lens_out.zero_() if q.shape[0] == 0: return out, lens_out if not q.is_cuda: @@ -433,11 +587,9 @@ def gluon_dsa_prefill_topk_fp8_gfx950( row_starts = row_starts.to(device=q.device, dtype=torch.int32).contiguous() row_ends = row_ends.to(device=q.device, dtype=torch.int32).contiguous() seq_len_sum = int(kv_workspace_slots.numel()) - candidate_lens = (row_ends - row_starts).clamp_min(0) - lens_out.copy_( - torch.minimum(candidate_lens, torch.full_like(candidate_lens, int(topk))) - ) if seq_len_sum == 0: + out.fill_(-1) + lens_out.zero_() return out, lens_out if max_logits_bytes is None: @@ -445,6 +597,8 @@ def gluon_dsa_prefill_topk_fp8_gfx950( else: max_query_rows = max(1, int(max_logits_bytes) // (max(seq_len_sum, 1) * 4)) block_n = 64 + select_warps = 8 + select_block = _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( @@ -472,11 +626,18 @@ def gluon_dsa_prefill_topk_fp8_gfx950( BLOCK_N=block_n, num_warps=4, ) - workspace_indices = _topk_with_padding(logits, int(topk)) - _workspace_topk_filter( - workspace_indices, + _dsa_prefill_select_topk_kernel[(end - start,)]( + logits, row_starts[start:end], row_ends[start:end], - out=out[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 From 9a280c2c0ff5e06828052c5a06251025fa348da3 Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Thu, 9 Jul 2026 16:40:40 +0000 Subject: [PATCH 06/14] Avoid Gluon DSA TopK long-row compile cliff Signed-off-by: Max Dawkins --- .../ops/attention/gluon/dsa_topk_gfx950.py | 37 ++++++++- .../test/ops/test_attention_dsa.py | 77 +++++++++++++++++++ 2 files changed, 113 insertions(+), 1 deletion(-) 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 index bf5d94749..8870894af 100644 --- 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 @@ -23,6 +23,11 @@ from __future__ import annotations import torch +from tokenspeed_kernel.ops.attention.triton.dsa_topk import ( + _RADIX_TOPK_MIN_COLS, + _topk_with_padding as _triton_topk_with_padding, + local_topk_to_global_slots as _triton_local_topk_to_global_slots, +) from tokenspeed_kernel_amd._triton import gl, gluon, triton __all__ = [ @@ -432,6 +437,10 @@ def _validate_topk(topk: int) -> None: raise ValueError(f"DSA Gluon top-k requires power-of-two topk, got {topk}") +def _use_triton_radix_topk(cols: int) -> bool: + return int(cols) >= _RADIX_TOPK_MIN_COLS + + def gluon_dsa_decode_topk_fp8_gfx950( q: torch.Tensor, weights: torch.Tensor, @@ -512,6 +521,17 @@ def gluon_dsa_decode_topk_fp8_gfx950( BLOCK_N=block_n, num_warps=4, ) + if _use_triton_radix_topk(max_seq_len): + local_topk_offsets = _triton_topk_with_padding(logits, topk) + return _triton_local_topk_to_global_slots( + local_topk_offsets=local_topk_offsets, + block_table=block_table, + block_size=int(page_size), + seq_lens=seq_lens, + out=out, + lens_out=lens_out, + ) + select_warps = 8 select_block = _next_power_of_2(max(max_seq_len, topk)) _dsa_decode_select_topk_kernel[(q.shape[0],)]( @@ -598,7 +618,14 @@ def gluon_dsa_prefill_topk_fp8_gfx950( max_query_rows = max(1, int(max_logits_bytes) // (max(seq_len_sum, 1) * 4)) block_n = 64 select_warps = 8 - select_block = _next_power_of_2(max(seq_len_sum, topk)) + use_triton_radix_topk = _use_triton_radix_topk(seq_len_sum) + if use_triton_radix_topk: + candidate_lens = (row_ends - row_starts).clamp_min(0) + lens_out.copy_( + torch.minimum(candidate_lens, torch.full_like(candidate_lens, topk)) + ) + else: + select_block = _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( @@ -626,6 +653,14 @@ def gluon_dsa_prefill_topk_fp8_gfx950( BLOCK_N=block_n, num_warps=4, ) + if use_triton_radix_topk: + workspace_indices = _triton_topk_with_padding(logits, topk) + valid = (workspace_indices >= row_starts[start:end, None]) & ( + workspace_indices < row_ends[start:end, None] + ) + out[start:end].copy_(torch.where(valid, workspace_indices, -1)) + continue + _dsa_prefill_select_topk_kernel[(end - start,)]( logits, row_starts[start:end], diff --git a/tokenspeed-kernel/test/ops/test_attention_dsa.py b/tokenspeed-kernel/test/ops/test_attention_dsa.py index 270a14e6c..6e260a96f 100644 --- a/tokenspeed-kernel/test/ops/test_attention_dsa.py +++ b/tokenspeed-kernel/test/ops/test_attention_dsa.py @@ -768,6 +768,83 @@ def test_dsa_prefill_topk_fp8_glm52_cases( _assert_topk_matches(workspace_indices, topk_lens, expected_indices, expected_lens) +def test_dsa_decode_topk_gluon_long_row_uses_radix_path( + device: str, + require, +) -> None: + require("attention", "dsa_decode_topk", "gluon", torch.bfloat16, "q") + + 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 = dsa_decode_topk( + q, + weights, + seq_lens, + block_table, + page_size=page_size, + topk=topk, + softmax_scale=head_dim**-0.5, + index_k_cache=packed_index_k, + solution="gluon", + ) + + 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( + device: str, + require, +) -> None: + require("attention", "dsa_prefill_topk", "gluon", torch.bfloat16, "q") + + 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 = dsa_prefill_topk( + 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, + solution="gluon", + ) + + 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 test_dsa_plan_triton(device: str) -> None: # The triton decode kernel derives its own causal bounds and ignores the # plan, so triton_dsa_plan is a no-op returning an opaque, non-None From ff356fed1d91c97ba59608ddec7318a4823394c4 Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Thu, 9 Jul 2026 17:36:16 +0000 Subject: [PATCH 07/14] Route long-row Gluon DSA TopK through Triton Signed-off-by: Max Dawkins --- .../ops/attention/gluon/dsa_topk_gfx950.py | 61 ++++---- .../test/ops/test_attention_dsa.py | 134 ++++++++++++++++++ 2 files changed, 166 insertions(+), 29 deletions(-) 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 index 8870894af..698397847 100644 --- 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 @@ -25,8 +25,8 @@ import torch from tokenspeed_kernel.ops.attention.triton.dsa_topk import ( _RADIX_TOPK_MIN_COLS, - _topk_with_padding as _triton_topk_with_padding, - local_topk_to_global_slots as _triton_local_topk_to_global_slots, + dsa_decode_topk_fp8 as _triton_dsa_decode_topk_fp8, + dsa_prefill_topk_fp8 as _triton_dsa_prefill_topk_fp8, ) from tokenspeed_kernel_amd._triton import gl, gluon, triton @@ -493,6 +493,20 @@ def gluon_dsa_decode_topk_fp8_gfx950( seq_lens = seq_lens.to(device=q.device, dtype=torch.int32).contiguous() block_table = block_table.to(device=q.device, dtype=torch.int32).contiguous() max_seq_len = int(block_table.shape[1]) * int(page_size) + if _use_triton_radix_topk(max_seq_len): + return _triton_dsa_decode_topk_fp8( + q, + index_k_cache, + weights, + seq_lens, + block_table, + page_size=int(page_size), + topk=topk, + softmax_scale=float(softmax_scale), + out=out, + lens_out=lens_out, + ) + if out is None: out = torch.empty((q.shape[0], topk), dtype=torch.int32, device=q.device) if lens_out is None: @@ -521,17 +535,6 @@ def gluon_dsa_decode_topk_fp8_gfx950( BLOCK_N=block_n, num_warps=4, ) - if _use_triton_radix_topk(max_seq_len): - local_topk_offsets = _triton_topk_with_padding(logits, topk) - return _triton_local_topk_to_global_slots( - local_topk_offsets=local_topk_offsets, - block_table=block_table, - block_size=int(page_size), - seq_lens=seq_lens, - out=out, - lens_out=lens_out, - ) - select_warps = 8 select_block = _next_power_of_2(max(max_seq_len, topk)) _dsa_decode_select_topk_kernel[(q.shape[0],)]( @@ -611,6 +614,21 @@ def gluon_dsa_prefill_topk_fp8_gfx950( out.fill_(-1) lens_out.zero_() return out, lens_out + if _use_triton_radix_topk(seq_len_sum): + return _triton_dsa_prefill_topk_fp8( + q, + index_k_cache, + weights, + kv_workspace_slots, + row_starts, + row_ends, + topk=topk, + softmax_scale=float(softmax_scale), + page_size=int(page_size), + max_logits_bytes=max_logits_bytes, + out=out, + lens_out=lens_out, + ) if max_logits_bytes is None: max_query_rows = q.shape[0] @@ -618,14 +636,7 @@ def gluon_dsa_prefill_topk_fp8_gfx950( max_query_rows = max(1, int(max_logits_bytes) // (max(seq_len_sum, 1) * 4)) block_n = 64 select_warps = 8 - use_triton_radix_topk = _use_triton_radix_topk(seq_len_sum) - if use_triton_radix_topk: - candidate_lens = (row_ends - row_starts).clamp_min(0) - lens_out.copy_( - torch.minimum(candidate_lens, torch.full_like(candidate_lens, topk)) - ) - else: - select_block = _next_power_of_2(max(seq_len_sum, topk)) + select_block = _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( @@ -653,14 +664,6 @@ def gluon_dsa_prefill_topk_fp8_gfx950( BLOCK_N=block_n, num_warps=4, ) - if use_triton_radix_topk: - workspace_indices = _triton_topk_with_padding(logits, topk) - valid = (workspace_indices >= row_starts[start:end, None]) & ( - workspace_indices < row_ends[start:end, None] - ) - out[start:end].copy_(torch.where(valid, workspace_indices, -1)) - continue - _dsa_prefill_select_topk_kernel[(end - start,)]( logits, row_starts[start:end], diff --git a/tokenspeed-kernel/test/ops/test_attention_dsa.py b/tokenspeed-kernel/test/ops/test_attention_dsa.py index 6e260a96f..c00a3b2b9 100644 --- a/tokenspeed-kernel/test/ops/test_attention_dsa.py +++ b/tokenspeed-kernel/test/ops/test_attention_dsa.py @@ -845,6 +845,140 @@ def test_dsa_prefill_topk_gluon_long_row_uses_radix_path( ) +def test_dsa_decode_topk_gluon_long_row_routes_full_triton( + device: str, + require, + monkeypatch: pytest.MonkeyPatch, +) -> None: + require("attention", "dsa_decode_topk", "gluon", torch.bfloat16, "q") + import tokenspeed_kernel_amd.ops.attention.gluon.dsa_topk_gfx950 as gluon_topk + + page_size = 64 + seq_len = 65536 + topk = 2048 + head_dim = 128 + q = torch.empty((1, 1, head_dim), device=device, dtype=torch.bfloat16) + weights = torch.empty((1, 1), device=device, dtype=torch.float32) + packed_index_k = torch.empty( + (seq_len, head_dim + 4), device=device, dtype=torch.uint8 + ) + 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) + out = torch.empty((1, topk), device=device, dtype=torch.int32) + lens_out = torch.empty((1,), device=device, dtype=torch.int32) + called = False + + def fake_triton_decode( + q_arg: torch.Tensor, + index_k_cache_arg: torch.Tensor, + weights_arg: torch.Tensor, + seq_lens_arg: torch.Tensor, + block_table_arg: torch.Tensor, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + nonlocal called + called = True + assert q_arg is q + assert index_k_cache_arg is packed_index_k + assert weights_arg is weights + assert seq_lens_arg.data_ptr() == seq_lens.data_ptr() + assert block_table_arg.data_ptr() == block_table.data_ptr() + assert kwargs["page_size"] == page_size + assert kwargs["topk"] == topk + assert kwargs["out"] is out + assert kwargs["lens_out"] is lens_out + return out, lens_out + + monkeypatch.setattr(gluon_topk, "_triton_dsa_decode_topk_fp8", fake_triton_decode) + + actual_out, actual_lens = gluon_topk.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, + out=out, + lens_out=lens_out, + ) + + assert called + assert actual_out is out + assert actual_lens is lens_out + + +def test_dsa_prefill_topk_gluon_long_row_routes_full_triton( + device: str, + require, + monkeypatch: pytest.MonkeyPatch, +) -> None: + require("attention", "dsa_prefill_topk", "gluon", torch.bfloat16, "q") + import tokenspeed_kernel_amd.ops.attention.gluon.dsa_topk_gfx950 as gluon_topk + + page_size = 64 + seq_len = 65536 + topk = 2048 + head_dim = 128 + q = torch.empty((1, 1, head_dim), device=device, dtype=torch.bfloat16) + weights = torch.empty((1, 1), device=device, dtype=torch.float32) + packed_index_k = torch.empty( + (seq_len, head_dim + 4), device=device, dtype=torch.uint8 + ) + 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) + out = torch.empty((1, topk), device=device, dtype=torch.int32) + lens_out = torch.empty((1,), device=device, dtype=torch.int32) + called = False + + def fake_triton_prefill( + q_arg: torch.Tensor, + index_k_cache_arg: torch.Tensor, + weights_arg: torch.Tensor, + kv_workspace_slots_arg: torch.Tensor, + row_starts_arg: torch.Tensor, + row_ends_arg: torch.Tensor, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + nonlocal called + called = True + assert q_arg is q + assert index_k_cache_arg is packed_index_k + assert weights_arg is weights + assert kv_workspace_slots_arg.data_ptr() == kv_workspace_slots.data_ptr() + assert row_starts_arg.data_ptr() == row_starts.data_ptr() + assert row_ends_arg.data_ptr() == row_ends.data_ptr() + assert kwargs["page_size"] == page_size + assert kwargs["topk"] == topk + assert kwargs["out"] is out + assert kwargs["lens_out"] is lens_out + return out, lens_out + + monkeypatch.setattr(gluon_topk, "_triton_dsa_prefill_topk_fp8", fake_triton_prefill) + + actual_out, actual_lens = gluon_topk.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, + out=out, + lens_out=lens_out, + ) + + assert called + assert actual_out is out + assert actual_lens is lens_out + + def test_dsa_plan_triton(device: str) -> None: # The triton decode kernel derives its own causal bounds and ignores the # plan, so triton_dsa_plan is a no-op returning an opaque, non-None From f82285a99072fc3232c214092ebec00bca7fb445 Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Thu, 9 Jul 2026 18:15:55 +0000 Subject: [PATCH 08/14] Keep long-row DSA TopK in Gluon Signed-off-by: Max Dawkins --- .../ops/attention/gluon/dsa_topk_gfx950.py | 591 +++++++++++++++--- .../test/ops/test_attention_dsa.py | 134 ---- 2 files changed, 521 insertions(+), 204 deletions(-) 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 index 698397847..36f3c5989 100644 --- 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 @@ -23,13 +23,11 @@ from __future__ import annotations import torch -from tokenspeed_kernel.ops.attention.triton.dsa_topk import ( - _RADIX_TOPK_MIN_COLS, - dsa_decode_topk_fp8 as _triton_dsa_decode_topk_fp8, - dsa_prefill_topk_fp8 as _triton_dsa_prefill_topk_fp8, -) from tokenspeed_kernel_amd._triton import gl, gluon, triton +_RADIX_TOPK_MIN_COLS = 65536 +_RADIX_TOPK_BLOCK_N = 4096 + __all__ = [ "gluon_dsa_decode_topk_fp8_gfx950", "gluon_dsa_prefill_topk_fp8_gfx950", @@ -45,6 +43,15 @@ def _vector_layout( return gl.BlockedLayout([LOAD_ELEMS], [64], [NUM_WARPS], [0]) +@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 _fp32_to_ordered_key(x): bits = x.to(gl.uint32, bitcast=True) @@ -243,18 +250,22 @@ def _dsa_decode_logits_fp8_kernel( 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) - offsets = block_id * BLOCK_N + gl.arange( - 0, BLOCK_N, layout=_vector_layout(BLOCK_N, gl.num_warps(), 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) seq_len = gl.load(seq_lens + token).to(gl.int32) valid = (offsets < seq_len) & (offsets < max_seq_len) block_idx = offsets // page_size block_offset = offsets - block_idx * page_size - page = gl.load( - block_table + token * block_table_stride + block_idx, + page = gl.amd.cdna4.buffer_load( + ptr=block_table, + offsets=(token * block_table_stride + block_idx).to(gl.int32), mask=valid, other=0, ).to(gl.int64) @@ -269,7 +280,7 @@ def _dsa_decode_logits_fp8_kernel( [BLOCK_N], value=0.0, dtype=gl.float32, - layout=_vector_layout(BLOCK_N, gl.num_warps(), 1), + layout=row_layout, ) for head in gl.static_range(0, num_heads): @@ -278,23 +289,29 @@ def _dsa_decode_logits_fp8_kernel( [BLOCK_N], value=0.0, dtype=gl.float32, - layout=_vector_layout(BLOCK_N, gl.num_warps(), 1), + layout=row_layout, ) - for dim in gl.static_range(0, head_dim): - q_val = gl.load(q + (token * num_heads + head) * head_dim + dim).to( - gl.float32 - ) - k_val = gl.load( - index_k_fp8 + fp8_base + dim, - mask=valid, + 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_scale = gl.load( - index_k_scale + scale_base + dim // 128, + 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 += k_val * k_scale * q_val + head_score += gl.sum(k_vals * k_scale[:, None] * q_vals[None, :], axis=1) scores += head_score * head_weight scores *= softmax_scale @@ -325,16 +342,24 @@ def _dsa_prefill_logits_fp8_kernel( 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) - offsets = block_id * BLOCK_N + gl.arange( - 0, BLOCK_N, layout=_vector_layout(BLOCK_N, gl.num_warps(), 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.load(kv_workspace_slots + offsets, mask=offsets < seq_len_sum, other=0) + 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 @@ -348,7 +373,7 @@ def _dsa_prefill_logits_fp8_kernel( [BLOCK_N], value=0.0, dtype=gl.float32, - layout=_vector_layout(BLOCK_N, gl.num_warps(), 1), + layout=row_layout, ) for head in gl.static_range(0, num_heads): @@ -357,23 +382,29 @@ def _dsa_prefill_logits_fp8_kernel( [BLOCK_N], value=0.0, dtype=gl.float32, - layout=_vector_layout(BLOCK_N, gl.num_warps(), 1), + layout=row_layout, ) - for dim in gl.static_range(0, head_dim): - q_val = gl.load(q + (token * num_heads + head) * head_dim + dim).to( - gl.float32 - ) - k_val = gl.load( - index_k_fp8 + fp8_base + dim, - mask=valid, + 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_scale = gl.load( - index_k_scale + scale_base + dim // 128, + 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 += k_val * k_scale * q_val + head_score += gl.sum(k_vals * k_scale[:, None] * q_vals[None, :], axis=1) scores += head_score * head_weight scores *= softmax_scale @@ -383,6 +414,266 @@ def _dsa_prefill_logits_fp8_kernel( ) +@gluon.jit +def _dsa_decode_radix_init_kernel( + seq_lens, + 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) + seq_len = gl.load(seq_lens + row).to(gl.int32) + 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 = gl.full((), 0, dtype=gl.int32) + selected = gl.full((), 0, dtype=gl.uint32) + selected_remaining = kth + found = gl.full((), 0, dtype=gl.int32) + + 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, + 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) + seq_len = gl.load(seq_lens + row).to(gl.int32) + 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 + row * 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 _check_packed_fp8_inputs( q: torch.Tensor, index_k_cache: torch.Tensor, @@ -437,10 +728,175 @@ def _validate_topk(topk: int) -> None: raise ValueError(f"DSA Gluon top-k requires power-of-two topk, got {topk}") -def _use_triton_radix_topk(cols: int) -> bool: +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 = _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) + 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, + 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, + 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, + 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, @@ -493,20 +949,6 @@ def gluon_dsa_decode_topk_fp8_gfx950( seq_lens = seq_lens.to(device=q.device, dtype=torch.int32).contiguous() block_table = block_table.to(device=q.device, dtype=torch.int32).contiguous() max_seq_len = int(block_table.shape[1]) * int(page_size) - if _use_triton_radix_topk(max_seq_len): - return _triton_dsa_decode_topk_fp8( - q, - index_k_cache, - weights, - seq_lens, - block_table, - page_size=int(page_size), - topk=topk, - softmax_scale=float(softmax_scale), - out=out, - lens_out=lens_out, - ) - if out is None: out = torch.empty((q.shape[0], topk), dtype=torch.int32, device=q.device) if lens_out is None: @@ -514,7 +956,7 @@ def gluon_dsa_decode_topk_fp8_gfx950( logits = torch.empty( (q.shape[0], max_seq_len), dtype=torch.float32, device=q.device ) - block_n = 64 + 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), @@ -533,8 +975,20 @@ def gluon_dsa_decode_topk_fp8_gfx950( 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(max_seq_len): + return _dsa_decode_radix_topk_slots( + logits, + block_table, + seq_lens, + page_size=int(page_size), + topk=topk, + out=out, + lens_out=lens_out, + ) + select_warps = 8 select_block = _next_power_of_2(max(max_seq_len, topk)) _dsa_decode_select_topk_kernel[(q.shape[0],)]( @@ -614,27 +1068,12 @@ def gluon_dsa_prefill_topk_fp8_gfx950( out.fill_(-1) lens_out.zero_() return out, lens_out - if _use_triton_radix_topk(seq_len_sum): - return _triton_dsa_prefill_topk_fp8( - q, - index_k_cache, - weights, - kv_workspace_slots, - row_starts, - row_ends, - topk=topk, - softmax_scale=float(softmax_scale), - page_size=int(page_size), - max_logits_bytes=max_logits_bytes, - out=out, - lens_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 = 64 + block_n = 32 select_warps = 8 select_block = _next_power_of_2(max(seq_len_sum, topk)) for start in range(0, q.shape[0], max_query_rows): @@ -662,8 +1101,20 @@ def gluon_dsa_prefill_topk_fp8_gfx950( 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], diff --git a/tokenspeed-kernel/test/ops/test_attention_dsa.py b/tokenspeed-kernel/test/ops/test_attention_dsa.py index c00a3b2b9..6e260a96f 100644 --- a/tokenspeed-kernel/test/ops/test_attention_dsa.py +++ b/tokenspeed-kernel/test/ops/test_attention_dsa.py @@ -845,140 +845,6 @@ def test_dsa_prefill_topk_gluon_long_row_uses_radix_path( ) -def test_dsa_decode_topk_gluon_long_row_routes_full_triton( - device: str, - require, - monkeypatch: pytest.MonkeyPatch, -) -> None: - require("attention", "dsa_decode_topk", "gluon", torch.bfloat16, "q") - import tokenspeed_kernel_amd.ops.attention.gluon.dsa_topk_gfx950 as gluon_topk - - page_size = 64 - seq_len = 65536 - topk = 2048 - head_dim = 128 - q = torch.empty((1, 1, head_dim), device=device, dtype=torch.bfloat16) - weights = torch.empty((1, 1), device=device, dtype=torch.float32) - packed_index_k = torch.empty( - (seq_len, head_dim + 4), device=device, dtype=torch.uint8 - ) - 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) - out = torch.empty((1, topk), device=device, dtype=torch.int32) - lens_out = torch.empty((1,), device=device, dtype=torch.int32) - called = False - - def fake_triton_decode( - q_arg: torch.Tensor, - index_k_cache_arg: torch.Tensor, - weights_arg: torch.Tensor, - seq_lens_arg: torch.Tensor, - block_table_arg: torch.Tensor, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - nonlocal called - called = True - assert q_arg is q - assert index_k_cache_arg is packed_index_k - assert weights_arg is weights - assert seq_lens_arg.data_ptr() == seq_lens.data_ptr() - assert block_table_arg.data_ptr() == block_table.data_ptr() - assert kwargs["page_size"] == page_size - assert kwargs["topk"] == topk - assert kwargs["out"] is out - assert kwargs["lens_out"] is lens_out - return out, lens_out - - monkeypatch.setattr(gluon_topk, "_triton_dsa_decode_topk_fp8", fake_triton_decode) - - actual_out, actual_lens = gluon_topk.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, - out=out, - lens_out=lens_out, - ) - - assert called - assert actual_out is out - assert actual_lens is lens_out - - -def test_dsa_prefill_topk_gluon_long_row_routes_full_triton( - device: str, - require, - monkeypatch: pytest.MonkeyPatch, -) -> None: - require("attention", "dsa_prefill_topk", "gluon", torch.bfloat16, "q") - import tokenspeed_kernel_amd.ops.attention.gluon.dsa_topk_gfx950 as gluon_topk - - page_size = 64 - seq_len = 65536 - topk = 2048 - head_dim = 128 - q = torch.empty((1, 1, head_dim), device=device, dtype=torch.bfloat16) - weights = torch.empty((1, 1), device=device, dtype=torch.float32) - packed_index_k = torch.empty( - (seq_len, head_dim + 4), device=device, dtype=torch.uint8 - ) - 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) - out = torch.empty((1, topk), device=device, dtype=torch.int32) - lens_out = torch.empty((1,), device=device, dtype=torch.int32) - called = False - - def fake_triton_prefill( - q_arg: torch.Tensor, - index_k_cache_arg: torch.Tensor, - weights_arg: torch.Tensor, - kv_workspace_slots_arg: torch.Tensor, - row_starts_arg: torch.Tensor, - row_ends_arg: torch.Tensor, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - nonlocal called - called = True - assert q_arg is q - assert index_k_cache_arg is packed_index_k - assert weights_arg is weights - assert kv_workspace_slots_arg.data_ptr() == kv_workspace_slots.data_ptr() - assert row_starts_arg.data_ptr() == row_starts.data_ptr() - assert row_ends_arg.data_ptr() == row_ends.data_ptr() - assert kwargs["page_size"] == page_size - assert kwargs["topk"] == topk - assert kwargs["out"] is out - assert kwargs["lens_out"] is lens_out - return out, lens_out - - monkeypatch.setattr(gluon_topk, "_triton_dsa_prefill_topk_fp8", fake_triton_prefill) - - actual_out, actual_lens = gluon_topk.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, - out=out, - lens_out=lens_out, - ) - - assert called - assert actual_out is out - assert actual_lens is lens_out - - def test_dsa_plan_triton(device: str) -> None: # The triton decode kernel derives its own causal bounds and ignores the # plan, so triton_dsa_plan is a no-op returning an opaque, non-None From 277e511a1421cce53b17a8cf382df361c92f3ffe Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Thu, 9 Jul 2026 21:00:09 +0000 Subject: [PATCH 09/14] Generalize Gluon DSA attention support Signed-off-by: Max Dawkins --- .../ops/attention/gluon/dsa_gfx950.py | 71 ++++++++++++++----- .../ops/attention/gluon/__init__.py | 10 +-- .../test/ops/test_attention_dsa.py | 26 +++++-- .../test/test_kernel_api_selection.py | 4 +- 4 files changed, 82 insertions(+), 29 deletions(-) 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 index 4d87a83ae..4ad027c24 100644 --- 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 @@ -772,20 +772,29 @@ def _check_inputs( qk_rope_head_dim: int, page_size: int, ) -> None: - if q.dtype != torch.bfloat16: - raise TypeError(f"Gluon DSA supports BF16 q for this milestone, got {q.dtype}") + 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 != 192: + if qk_nope_head_dim not in (128, 192): raise ValueError( - f"Gluon DSA supports qk_nope_head_dim=192, got {qk_nope_head_dim}" + "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 kv_lora_rank != 512: - raise ValueError(f"Gluon DSA supports kv_lora_rank=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: @@ -794,6 +803,10 @@ def _check_inputs( 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 _run_dense_kv( q: torch.Tensor, kv_cache: torch.Tensor, @@ -805,7 +818,9 @@ def _run_dense_kv( qk_rope_head_dim: int, ) -> torch.Tensor: out = torch.empty( - (q.shape[0], q.shape[1], kv_lora_rank), dtype=q.dtype, device=q.device + (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]( @@ -844,7 +859,9 @@ def _run_dense_kv_scalar( ) -> 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=q.dtype, device=q.device + (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, @@ -878,7 +895,9 @@ def _run_packed_kv( ) -> torch.Tensor: row_bytes = int(packed_kv.shape[1]) out = torch.empty( - (q.shape[0], q.shape[1], kv_lora_rank), dtype=q.dtype, device=q.device + (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, @@ -943,15 +962,31 @@ def _run_dsa( qk_rope_head_dim=qk_rope_head_dim, ) elif kv_cache is not None: - result = _run_dense_kv( - q, - _flatten_dense_kv_cache(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, - ) + 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: 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 55c9efe7e..74bee8412 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py @@ -228,14 +228,15 @@ def gluon_dsa_prefill_topk_fp8_gfx950(*args, **kwargs): 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({192}), - "kv_lora_rank": frozenset({512}), + "qk_nope_head_dim": frozenset({128, 192}), + "kv_lora_rank": frozenset({128, 512}), "qk_rope_head_dim": frozenset({64}), "topk": frozenset({512, 1024, 2048}), "kv_cache_available": frozenset({False, True}), @@ -261,14 +262,15 @@ def gluon_dsa_decode_gfx950(*args, **kwargs): 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}), - "qk_nope_head_dim": frozenset({192}), - "kv_lora_rank": frozenset({512}), + "qk_nope_head_dim": frozenset({128, 192}), + "kv_lora_rank": frozenset({128, 512}), "qk_rope_head_dim": frozenset({64}), "topk": frozenset({512, 1024, 2048}), "kv_cache_available": frozenset({False, True}), diff --git a/tokenspeed-kernel/test/ops/test_attention_dsa.py b/tokenspeed-kernel/test/ops/test_attention_dsa.py index 6e260a96f..29ee6e7f7 100644 --- a/tokenspeed-kernel/test/ops/test_attention_dsa.py +++ b/tokenspeed-kernel/test/ops/test_attention_dsa.py @@ -1010,7 +1010,7 @@ def _assert_slots_visible( pytest.param("prefill", "dsa_prefill", id="prefill"), ], ) -@pytest.mark.parametrize("solution", ["triton"]) +@pytest.mark.parametrize("solution", ["triton", "gluon"]) @pytest.mark.parametrize( "q_dtype", [ @@ -1089,8 +1089,23 @@ def test_dsa_with_kvcache( pytest.param(torch.float8_e4m3fn, id="q_fp8"), ], ) -def test_dsa_decode_dense_kvcache(device: str, q_dtype: torch.dtype, require) -> None: - require("attention", "dsa_decode", "triton", q_dtype, "q") +@pytest.mark.parametrize( + "mode,api_name", + [ + pytest.param("decode", "dsa_decode", id="decode"), + pytest.param("prefill", "dsa_prefill", id="prefill"), + ], +) +@pytest.mark.parametrize("solution", ["triton", "gluon"]) +def test_dsa_dense_kvcache( + device: str, + q_dtype: torch.dtype, + mode: str, + api_name: str, + solution: str, + require, +) -> None: + require("attention", api_name, solution, q_dtype, "q") tokens = 3 num_heads = 2 @@ -1119,7 +1134,8 @@ def test_dsa_decode_dense_kvcache(device: str, q_dtype: torch.dtype, require) -> count = int(topk_lens[token].item()) topk_slots[token, :count] = torch.randperm(num_slots, device=device)[:count] - out = dsa_decode( + api = dsa_decode if mode == "decode" else dsa_prefill + out = api( q=q, kv_cache=kv_cache, sparse_kv_cache=None, @@ -1131,7 +1147,7 @@ def test_dsa_decode_dense_kvcache(device: str, q_dtype: torch.dtype, require) -> qk_rope_head_dim=qk_rope_head_dim, softmax_scale=softmax_scale, page_size=64, - solution="triton", + solution=solution, ) ref = _dsa_reference( diff --git a/tokenspeed-kernel/test/test_kernel_api_selection.py b/tokenspeed-kernel/test/test_kernel_api_selection.py index 60965e838..6d64f536d 100644 --- a/tokenspeed-kernel/test/test_kernel_api_selection.py +++ b/tokenspeed-kernel/test/test_kernel_api_selection.py @@ -1146,7 +1146,7 @@ def _case( "cdna4", "attention", "dsa_decode_topk", - "triton_dsa_decode_topk_fp8", + "gluon_dsa_decode_topk_fp8_gfx950", _attention_dsa_decode_topk, ), _case( @@ -1154,7 +1154,7 @@ def _case( "cdna4", "attention", "dsa_prefill_topk", - "triton_dsa_prefill_topk_fp8", + "gluon_dsa_prefill_topk_fp8_gfx950", _attention_dsa_prefill_topk, ), _case( From 3c0d82516472c651d04747d89cfdce1d2830725e Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Thu, 9 Jul 2026 21:11:33 +0000 Subject: [PATCH 10/14] Split Gluon DSA TopK scoring kernels Signed-off-by: Max Dawkins --- .../ops/attention/gluon/dsa_score_gfx950.py | 261 ++++++++++++++++++ .../ops/attention/gluon/dsa_topk_gfx950.py | 234 +--------------- 2 files changed, 266 insertions(+), 229 deletions(-) create mode 100644 tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_score_gfx950.py 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..58f4278fd --- /dev/null +++ b/tokenspeed-kernel-amd/python/tokenspeed_kernel_amd/ops/attention/gluon/dsa_score_gfx950.py @@ -0,0 +1,261 @@ +# 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, + 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) + seq_len = gl.load(seq_lens + token).to(gl.int32) + 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=(token * 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 index 36f3c5989..a1822aa69 100644 --- 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 @@ -24,6 +24,11 @@ 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 @@ -43,15 +48,6 @@ def _vector_layout( return gl.BlockedLayout([LOAD_ELEMS], [64], [NUM_WARPS], [0]) -@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 _fp32_to_ordered_key(x): bits = x.to(gl.uint32, bitcast=True) @@ -231,189 +227,6 @@ def _dsa_prefill_select_topk_kernel( gl.store(out + row * out_stride + selected_pos, offsets.to(gl.int32), mask=write) -@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, - 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) - seq_len = gl.load(seq_lens + token).to(gl.int32) - 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=(token * 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 - ) - - @gluon.jit def _dsa_decode_radix_init_kernel( seq_lens, @@ -674,43 +487,6 @@ def _dsa_prefill_radix_scatter_kernel( gl.store(out + row * out_stride + equal_pos, local, mask=equal_write) -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 - - def _next_power_of_2(value: int) -> int: if value <= 1: return 1 From 1b614645cf65bee6490573a8e4b007e858673d55 Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Fri, 10 Jul 2026 14:18:55 +0000 Subject: [PATCH 11/14] Move Gluon DSA tests to AMD package Signed-off-by: Max Dawkins --- .../ops/test_attention_dsa_gluon_gfx950.py | 963 ++++++++++++++++++ .../test/ops/test_attention_dsa.py | 777 +------------- 2 files changed, 968 insertions(+), 772 deletions(-) create mode 100644 tokenspeed-kernel-amd/test/ops/test_attention_dsa_gluon_gfx950.py 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..34cdbea9d --- /dev/null +++ b/tokenspeed-kernel-amd/test/ops/test_attention_dsa_gluon_gfx950.py @@ -0,0 +1,963 @@ +# 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.dsa_gfx950 import ( # noqa: E402 + 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) + + +@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, + ), +) + + +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, +) -> tuple[torch.Tensor, torch.Tensor]: + out = torch.full((q.shape[0], topk), -1, device=q.device, dtype=torch.int32) + lens = torch.minimum(seq_lens.to(torch.int32), torch.full_like(seq_lens, int(topk))) + for token in range(q.shape[0]): + seq_len = int(seq_lens[token].item()) + count = int(lens[token].item()) + if count == 0: + continue + offsets = torch.arange(seq_len, device=q.device, dtype=torch.long) + pages = block_table[token].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() + + +@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) + q = _randn_bf16( + (len(case.seq_lens), case.index_heads, head_dim), + device=device, + generator=gen, + ) + weights = _normal_weights( + (len(case.seq_lens), 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, + 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, + ) + + _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_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) + + +@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/test/ops/test_attention_dsa.py b/tokenspeed-kernel/test/ops/test_attention_dsa.py index 29ee6e7f7..50e19c6ef 100644 --- a/tokenspeed-kernel/test/ops/test_attention_dsa.py +++ b/tokenspeed-kernel/test/ops/test_attention_dsa.py @@ -21,8 +21,6 @@ from __future__ import annotations import math -from collections.abc import Sequence -from dataclasses import dataclass import pytest import torch @@ -40,246 +38,6 @@ torch.manual_seed(42) -@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, - ), -) - - def _pack_index_k_cache( index_k: torch.Tensor, page_size: int, @@ -316,177 +74,6 @@ def _pack_index_k_cache( 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, -) -> tuple[torch.Tensor, torch.Tensor]: - out = torch.full((q.shape[0], topk), -1, device=q.device, dtype=torch.int32) - lens = torch.minimum(seq_lens.to(torch.int32), torch.full_like(seq_lens, int(topk))) - for token in range(q.shape[0]): - seq_len = int(seq_lens[token].item()) - count = int(lens[token].item()) - if count == 0: - continue - offsets = torch.arange(seq_len, device=q.device, dtype=torch.long) - pages = block_table[token].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 test_dsa_decode_topk_fp8(device: str, require) -> None: require("attention", "dsa_decode_topk", "triton", torch.bfloat16, "q") @@ -647,204 +234,6 @@ def test_dsa_prefill_topk_fp8(device: str, require) -> None: assert (workspace_indices[0, int(expected_lens[0].item()) :] == -1).all() -@pytest.mark.parametrize("solution", ["triton", "gluon"]) -@pytest.mark.parametrize( - "case", - _GLM52_TOPK_DECODE_CASES, - ids=lambda case: case.name, -) -def test_dsa_decode_topk_fp8_glm52_cases( - device: str, - require, - solution: str, - case: _TopKDecodeCase, -) -> None: - require("attention", "dsa_decode_topk", solution, torch.bfloat16, "q") - - 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) - q = _randn_bf16( - (len(case.seq_lens), case.index_heads, head_dim), - device=device, - generator=gen, - ) - weights = _normal_weights( - (len(case.seq_lens), 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 = dsa_decode_topk( - q, - weights, - seq_lens, - block_table, - page_size=page_size, - topk=case.topk, - softmax_scale=softmax_scale, - q_len_per_req=case.q_len_per_req, - index_k_cache=packed_index_k, - solution=solution, - ) - 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, - ) - - _assert_topk_matches(topk_slots, topk_lens, expected_slots, expected_lens) - - -@pytest.mark.parametrize("solution", ["triton", "gluon"]) -@pytest.mark.parametrize( - "case", - _GLM52_TOPK_PREFILL_CASES, - ids=lambda case: case.name, -) -def test_dsa_prefill_topk_fp8_glm52_cases( - device: str, - require, - solution: str, - case: _TopKPrefillCase, -) -> None: - require("attention", "dsa_prefill_topk", solution, torch.bfloat16, "q") - - 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 = dsa_prefill_topk( - 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, - solution=solution, - ) - 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_decode_topk_gluon_long_row_uses_radix_path( - device: str, - require, -) -> None: - require("attention", "dsa_decode_topk", "gluon", torch.bfloat16, "q") - - 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 = dsa_decode_topk( - q, - weights, - seq_lens, - block_table, - page_size=page_size, - topk=topk, - softmax_scale=head_dim**-0.5, - index_k_cache=packed_index_k, - solution="gluon", - ) - - 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( - device: str, - require, -) -> None: - require("attention", "dsa_prefill_topk", "gluon", torch.bfloat16, "q") - - 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 = dsa_prefill_topk( - 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, - solution="gluon", - ) - - 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 test_dsa_plan_triton(device: str) -> None: # The triton decode kernel derives its own causal bounds and ignores the # plan, so triton_dsa_plan is a no-op returning an opaque, non-None @@ -928,12 +317,8 @@ def _dsa_reference( 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) @@ -943,66 +328,6 @@ def _dsa_reference( 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_name", [ @@ -1010,7 +335,7 @@ def _assert_slots_visible( pytest.param("prefill", "dsa_prefill", id="prefill"), ], ) -@pytest.mark.parametrize("solution", ["triton", "gluon"]) +@pytest.mark.parametrize("solution", ["triton"]) @pytest.mark.parametrize( "q_dtype", [ @@ -1089,23 +414,8 @@ def test_dsa_with_kvcache( pytest.param(torch.float8_e4m3fn, id="q_fp8"), ], ) -@pytest.mark.parametrize( - "mode,api_name", - [ - pytest.param("decode", "dsa_decode", id="decode"), - pytest.param("prefill", "dsa_prefill", id="prefill"), - ], -) -@pytest.mark.parametrize("solution", ["triton", "gluon"]) -def test_dsa_dense_kvcache( - device: str, - q_dtype: torch.dtype, - mode: str, - api_name: str, - solution: str, - require, -) -> None: - require("attention", api_name, solution, q_dtype, "q") +def test_dsa_decode_dense_kvcache(device: str, q_dtype: torch.dtype, require) -> None: + require("attention", "dsa_decode", "triton", q_dtype, "q") tokens = 3 num_heads = 2 @@ -1134,8 +444,7 @@ def test_dsa_dense_kvcache( count = int(topk_lens[token].item()) topk_slots[token, :count] = torch.randperm(num_slots, device=device)[:count] - api = dsa_decode if mode == "decode" else dsa_prefill - out = api( + out = dsa_decode( q=q, kv_cache=kv_cache, sparse_kv_cache=None, @@ -1147,7 +456,7 @@ def test_dsa_dense_kvcache( qk_rope_head_dim=qk_rope_head_dim, softmax_scale=softmax_scale, page_size=64, - solution=solution, + solution="triton", ) ref = _dsa_reference( @@ -1161,79 +470,3 @@ def test_dsa_dense_kvcache( 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, -) -@pytest.mark.parametrize("solution", ["triton", "gluon"]) -def test_dsa_glm52_selected_attention_cases( - device: str, - require, - solution: str, - case: _DSACase, -) -> None: - api_name = "dsa_decode" if case.mode == "decode" else "dsa_prefill" - require("attention", api_name, solution, torch.bfloat16, "q") - - 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, - "solution": solution, - } - if case.mode == "decode": - out = dsa_decode(q_len_per_req=case.q_len_per_req, **common_kwargs) - else: - out = dsa_prefill(**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) From eb17ddb7067db8521f1875b0f7372ed5efe992bf Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Fri, 10 Jul 2026 15:38:40 +0000 Subject: [PATCH 12/14] Select Gluon DSA by default in golden test Signed-off-by: Max Dawkins --- .../tokenspeed_kernel_amd/ops/attention/gluon/__init__.py | 6 +++--- .../tokenspeed_kernel/ops/attention/gluon/__init__.py | 6 +++--- tokenspeed-kernel/test/test_kernel_api_selection.py | 6 ++---- 3 files changed, 8 insertions(+), 10 deletions(-) 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 b95e849a9..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,9 +22,6 @@ from __future__ import annotations -from tokenspeed_kernel_amd.ops.attention.gluon.mha_decode_gfx950 import ( # noqa: F401 - gluon_mha_decode_gfx950, -) from tokenspeed_kernel_amd.ops.attention.gluon.dsa_gfx950 import ( # noqa: F401 gluon_dsa_decode_gfx950, gluon_dsa_prefill_gfx950, @@ -33,6 +30,9 @@ 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, +) from tokenspeed_kernel_amd.ops.attention.gluon.mha_extend_gfx950 import ( # noqa: F401 gluon_mha_extend_gfx950, ) 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 74bee8412..386ca622b 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py @@ -36,9 +36,6 @@ ) if current_platform().is_amd: - from tokenspeed_kernel_amd.ops.attention.gluon.mha_decode_gfx950 import ( - gluon_mha_decode_gfx950 as _decode_impl, - ) from tokenspeed_kernel_amd.ops.attention.gluon.dsa_gfx950 import ( gluon_dsa_decode_gfx950 as _dsa_decode_impl, ) @@ -51,6 +48,9 @@ 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, + ) from tokenspeed_kernel_amd.ops.attention.gluon.mha_extend_gfx950 import ( gluon_mha_extend_gfx950 as _extend_impl, ) diff --git a/tokenspeed-kernel/test/test_kernel_api_selection.py b/tokenspeed-kernel/test/test_kernel_api_selection.py index 6d64f536d..9090c79f5 100644 --- a/tokenspeed-kernel/test/test_kernel_api_selection.py +++ b/tokenspeed-kernel/test/test_kernel_api_selection.py @@ -499,7 +499,6 @@ def _attention_dsa_decode() -> object: qk_rope_head_dim=64, softmax_scale=1.0, page_size=64, - solution="triton", ) @@ -520,7 +519,6 @@ def _attention_dsa_prefill() -> object: qk_rope_head_dim=64, softmax_scale=1.0, page_size=64, - solution="triton", ) @@ -1130,7 +1128,7 @@ def _case( "cdna4", "attention", "dsa_decode", - "triton_dsa_decode", + "gluon_dsa_decode_gfx950", _attention_dsa_decode, ), _case( @@ -1138,7 +1136,7 @@ def _case( "cdna4", "attention", "dsa_prefill", - "triton_dsa_prefill", + "gluon_dsa_prefill_gfx950", _attention_dsa_prefill, ), _case( From be314f96284cfd16f51096286bcf3b899c121058 Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Tue, 14 Jul 2026 14:35:22 +0000 Subject: [PATCH 13/14] Fix Gluon DSA sparse prefill edge cases Signed-off-by: Max Dawkins --- .../ops/attention/gluon/dsa_gfx950.py | 26 +- .../ops/attention/gluon/dsa_score_gfx950.py | 9 +- .../ops/attention/gluon/dsa_topk_gfx950.py | 90 +++++-- .../ops/test_attention_dsa_gluon_gfx950.py | 253 +++++++++++++++++- .../ops/attention/gluon/__init__.py | 8 +- .../test/test_kernel_api_selection.py | 28 ++ 6 files changed, 376 insertions(+), 38 deletions(-) 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 index 4ad027c24..5996d9aed 100644 --- 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 @@ -30,6 +30,8 @@ "gluon_dsa_prefill_gfx950", ] +_REGISTERED_TOPK_WIDTHS = (512, 1024, 2048) + @gluon.constexpr_function def _value_layout( @@ -807,6 +809,22 @@ 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, @@ -935,6 +953,7 @@ def _run_dsa( page_size: int, k_scale: float, out: torch.Tensor | None, + max_seqlen_k: int, ) -> torch.Tensor: _check_inputs( q, @@ -948,6 +967,7 @@ def _run_dsa( 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. @@ -1014,7 +1034,7 @@ def gluon_dsa_decode_gfx950( return_lse: bool = False, out: torch.Tensor | None = None, ) -> torch.Tensor: - del max_seqlen_k, q_len_per_req + 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( @@ -1030,6 +1050,7 @@ def gluon_dsa_decode_gfx950( page_size=page_size, k_scale=k_scale, out=out, + max_seqlen_k=max_seqlen_k, ) @@ -1051,7 +1072,7 @@ def gluon_dsa_prefill_gfx950( return_lse: bool = False, out: torch.Tensor | None = None, ) -> torch.Tensor: - del max_seqlen_k, q_len_per_req + 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( @@ -1067,4 +1088,5 @@ def gluon_dsa_prefill_gfx950( 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 index 58f4278fd..3a63da043 100644 --- 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 @@ -59,6 +59,7 @@ def _dsa_decode_logits_fp8_kernel( 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, ): @@ -69,13 +70,17 @@ def _dsa_decode_logits_fp8_kernel( 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) - seq_len = gl.load(seq_lens + token).to(gl.int32) + 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=(token * block_table_stride + block_idx).to(gl.int32), + offsets=(req * block_table_stride + block_idx).to(gl.int32), mask=valid, other=0, ).to(gl.int64) 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 index a1822aa69..075baaf50 100644 --- 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 @@ -120,6 +120,7 @@ def _dsa_decode_select_topk_kernel( 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, @@ -129,7 +130,11 @@ def _dsa_decode_select_topk_kernel( 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) - seq_len = gl.load(seq_lens + row).to(gl.int32) + 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) @@ -140,7 +145,7 @@ def _dsa_decode_select_topk_kernel( block_idx = local // page_size block_offset = local - block_idx * page_size page = gl.load( - block_table + row * block_table_stride + block_idx, + block_table + req * block_table_stride + block_idx, mask=valid_top & (block_idx < block_table_cols), other=0, ).to(gl.int32) @@ -160,20 +165,26 @@ def _dsa_decode_select_topk_kernel( ) threshold = _find_topk_threshold_key(values, valid, topk, BLOCK_N, layout) keys = _fp32_to_ordered_key(values) - selected = valid & (keys >= threshold) - selected_i32 = selected.to(gl.int32) - selected_pos = gl.associative_scan(selected_i32, 0, _topk_add) - 1 - write = selected & (selected_pos < topk) + 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 + row * block_table_stride + block_idx, - mask=write & (block_idx < block_table_cols), + 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 + selected_pos, slots, mask=write) + 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 @@ -220,11 +231,18 @@ def _dsa_prefill_select_topk_kernel( ) threshold = _find_topk_threshold_key(values, valid, topk, BLOCK_N, layout) keys = _fp32_to_ordered_key(values) - selected = valid & (keys >= threshold) - selected_i32 = selected.to(gl.int32) - selected_pos = gl.associative_scan(selected_i32, 0, _topk_add) - 1 - write = selected & (selected_pos < topk) - gl.store(out + row * out_stride + selected_pos, offsets.to(gl.int32), mask=write) + 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 @@ -237,12 +255,17 @@ def _dsa_decode_radix_init_kernel( 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) - seq_len = gl.load(seq_lens + row).to(gl.int32) + 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) @@ -377,6 +400,7 @@ def _dsa_decode_radix_scatter_slots_kernel( 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, ): @@ -384,7 +408,11 @@ def _dsa_decode_radix_scatter_slots_kernel( 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) - seq_len = gl.load(seq_lens + row).to(gl.int32) + 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, @@ -420,7 +448,7 @@ def _dsa_decode_radix_scatter_slots_kernel( block_idx = offsets.to(gl.int32) // page_size block_offset = offsets.to(gl.int32) - block_idx * page_size page = gl.load( - block_table + row * block_table_stride + block_idx, + block_table + req * block_table_stride + block_idx, mask=(greater_write | equal_write) & (block_idx < block_table_cols), other=0, ).to(gl.int32) @@ -567,6 +595,7 @@ def _dsa_decode_radix_topk_slots( *, page_size: int, topk: int, + q_len_per_req: int, out: torch.Tensor, lens_out: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -583,6 +612,7 @@ def _dsa_decode_radix_topk_slots( counters, out.stride(0), topk=topk, + q_len_per_req=q_len_per_req, TOPK_LOAD_ELEMS=_load_elems(topk, 8), num_warps=8, ) @@ -611,6 +641,7 @@ def _dsa_decode_radix_topk_slots( 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, @@ -684,24 +715,34 @@ def gluon_dsa_decode_topk_fp8_gfx950( 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, q_len_per_req + 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 or seq_lens.numel() != q.shape[0]: + 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( - "seq_lens must be [tokens], got " - f"{tuple(seq_lens.shape)} for q={tuple(q.shape)}" + "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] < q.shape[0]: + if block_table.dim() != 2 or block_table.shape[0] < seq_lens.numel(): raise ValueError( - "block_table must have at least one row per token, got " + "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: @@ -750,6 +791,7 @@ def gluon_dsa_decode_topk_fp8_gfx950( 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, @@ -761,6 +803,7 @@ def gluon_dsa_decode_topk_fp8_gfx950( seq_lens, page_size=int(page_size), topk=topk, + q_len_per_req=q_len_per_req, out=out, lens_out=lens_out, ) @@ -779,6 +822,7 @@ def gluon_dsa_decode_topk_fp8_gfx950( 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), 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 index 34cdbea9d..c984ff85a 100644 --- a/tokenspeed-kernel-amd/test/ops/test_attention_dsa_gluon_gfx950.py +++ b/tokenspeed-kernel-amd/test/ops/test_attention_dsa_gluon_gfx950.py @@ -21,7 +21,9 @@ def _is_gfx950() -> bool: 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, ) @@ -33,6 +35,24 @@ def _is_gfx950() -> bool: 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 @@ -270,6 +290,17 @@ class _DSACase: 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, + ), ) @@ -404,16 +435,22 @@ def _reference_decode_topk( 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.minimum(seq_lens.to(torch.int32), torch.full_like(seq_lens, int(topk))) + lens = torch.empty((q.shape[0],), device=q.device, dtype=torch.int32) for token in range(q.shape[0]): - seq_len = int(seq_lens[token].item()) - count = int(lens[token].item()) + 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[token].long().index_select(0, offsets // page_size) + 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], @@ -489,14 +526,13 @@ def test_dsa_decode_topk_fp8_glm52_cases(case: _TopKDecodeCase) -> None: 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( - (len(case.seq_lens), case.index_heads, head_dim), + (tokens, case.index_heads, head_dim), device=device, generator=gen, ) - weights = _normal_weights( - (len(case.seq_lens), case.index_heads), 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, @@ -511,6 +547,7 @@ def test_dsa_decode_topk_fp8_glm52_cases(case: _TopKDecodeCase) -> None: 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, ) @@ -523,6 +560,7 @@ def test_dsa_decode_topk_fp8_glm52_cases(case: _TopKDecodeCase) -> None: 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) @@ -582,6 +620,153 @@ def test_dsa_prefill_topk_fp8_glm52_cases(case: _TopKPrefillCase) -> None: _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 @@ -896,6 +1081,58 @@ def test_dsa_dense_kvcache(mode: str, api, q_dtype: torch.dtype) -> None: 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, 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 386ca622b..b286116b4 100644 --- a/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py +++ b/tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/gluon/__init__.py @@ -36,6 +36,9 @@ ) 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, ) @@ -238,7 +241,7 @@ def gluon_dsa_prefill_topk_fp8_gfx950(*args, **kwargs): "qk_nope_head_dim": frozenset({128, 192}), "kv_lora_rank": frozenset({128, 512}), "qk_rope_head_dim": frozenset({64}), - "topk": frozenset({512, 1024, 2048}), + "topk": _DSA_FULL_TOPK_WIDTHS, "kv_cache_available": frozenset({False, True}), "sparse_kv_cache_available": frozenset({False, True}), "topk_layout": frozenset({"global_slots"}), @@ -262,7 +265,6 @@ def gluon_dsa_decode_gfx950(*args, **kwargs): signatures=frozenset( { format_signature(q=dense_tensor_format(torch.bfloat16)), - format_signature(q=dense_tensor_format(torch.float8_e4m3fn)), } ), priority=Priority.SPECIALIZED, @@ -272,7 +274,7 @@ def gluon_dsa_decode_gfx950(*args, **kwargs): "qk_nope_head_dim": frozenset({128, 192}), "kv_lora_rank": frozenset({128, 512}), "qk_rope_head_dim": frozenset({64}), - "topk": frozenset({512, 1024, 2048}), + "topk": _DSA_PREFILL_TOPK_WIDTHS, "kv_cache_available": frozenset({False, True}), "sparse_kv_cache_available": frozenset({False, True}), "topk_layout": frozenset({"global_slots"}), diff --git a/tokenspeed-kernel/test/test_kernel_api_selection.py b/tokenspeed-kernel/test/test_kernel_api_selection.py index 9090c79f5..500648d7c 100644 --- a/tokenspeed-kernel/test/test_kernel_api_selection.py +++ b/tokenspeed-kernel/test/test_kernel_api_selection.py @@ -522,6 +522,26 @@ def _attention_dsa_prefill() -> object: ) +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, + ) + + def _attention_dsa_decode_topk() -> object: q = torch.empty((2, 2, 128), dtype=torch.bfloat16) weights = torch.empty((2, 2), dtype=torch.float32) @@ -1139,6 +1159,14 @@ def _case( "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", From 5fe5c416987e30391783b147a7438c72f3ab03e2 Mon Sep 17 00:00:00 2001 From: Max Dawkins Date: Tue, 14 Jul 2026 14:37:01 +0000 Subject: [PATCH 14/14] Address Gluon DSA TopK review comments Signed-off-by: Max Dawkins --- .../ops/attention/gluon/dsa_topk_gfx950.py | 76 ++++++++++--------- .../ops/test_attention_dsa_gluon_gfx950.py | 73 ++++++++++++++++++ 2 files changed, 113 insertions(+), 36 deletions(-) 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 index 075baaf50..8ed370646 100644 --- 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 @@ -69,19 +69,20 @@ def _find_topk_threshold_key( layout: gl.constexpr, ): keys = _fp32_to_ordered_key(values) - prefix = gl.full((), 0, dtype=gl.uint32) - remaining = gl.full((), topk, dtype=gl.int32) + 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 = gl.full((), 0, dtype=gl.int32) - selected = gl.full((), 0, dtype=gl.uint32) + cumulative = 0 + selected = 0 selected_remaining = remaining - found = gl.full((), 0, dtype=gl.int32) + found = 0 for bucket_id in gl.static_range(15, -1, -1): in_bucket = prefix_match & (bucket == bucket_id) @@ -360,10 +361,10 @@ def _dsa_radix_update_kernel( tile_mask = tile_offsets < hist_tiles row_hist = hist + row * hist_tiles * 16 kth = gl.load(remaining + row).to(gl.int32) - cumulative = gl.full((), 0, dtype=gl.int32) - selected = gl.full((), 0, dtype=gl.uint32) + cumulative = 0 + selected = 0 selected_remaining = kth - found = gl.full((), 0, dtype=gl.int32) + found = 0 for bucket_desc in gl.static_range(0, 16): bucket_id = 15 - bucket_desc @@ -515,16 +516,24 @@ def _dsa_prefill_radix_scatter_kernel( gl.store(out + row * out_stride + equal_pos, local, mask=equal_write) -def _next_power_of_2(value: int) -> int: - if value <= 1: - return 1 - return 1 << (int(value) - 1).bit_length() - - 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}") @@ -547,7 +556,7 @@ def _dsa_radix_scratch( 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 = _next_power_of_2(tiles) + block_tiles = triton.next_power_of_2(tiles) return hist, prefixes, remaining, counters, tiles, block_tiles @@ -564,6 +573,7 @@ def _run_radix_prefix_passes( ) -> 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, @@ -757,14 +767,11 @@ def gluon_dsa_decode_topk_fp8_gfx950( else lens_out ) return empty_out, empty_lens - if not q.is_cuda: - raise RuntimeError("DSA Gluon FP8 decode top-k requires CUDA tensors") - - q = q.contiguous() - index_k_cache = index_k_cache.contiguous() - weights = weights.contiguous() - seq_lens = seq_lens.to(device=q.device, dtype=torch.int32).contiguous() - block_table = block_table.to(device=q.device, dtype=torch.int32).contiguous() + 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) @@ -809,7 +816,7 @@ def gluon_dsa_decode_topk_fp8_gfx950( ) select_warps = 8 - select_block = _next_power_of_2(max(max_seq_len, topk)) + select_block = triton.next_power_of_2(max(max_seq_len, topk)) _dsa_decode_select_topk_kernel[(q.shape[0],)]( logits, block_table, @@ -872,17 +879,14 @@ def gluon_dsa_prefill_topk_fp8_gfx950( lens_out = torch.empty((q.shape[0],), dtype=torch.int32, device=q.device) if q.shape[0] == 0: return out, lens_out - if not q.is_cuda: - raise RuntimeError("DSA Gluon FP8 prefill top-k requires CUDA tensors") - - q = q.contiguous() - index_k_cache = index_k_cache.contiguous() - weights = weights.contiguous() - kv_workspace_slots = kv_workspace_slots.to( - device=q.device, dtype=torch.int64 - ).contiguous() - row_starts = row_starts.to(device=q.device, dtype=torch.int32).contiguous() - row_ends = row_ends.to(device=q.device, dtype=torch.int32).contiguous() + 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) @@ -895,7 +899,7 @@ def gluon_dsa_prefill_topk_fp8_gfx950( max_query_rows = max(1, int(max_logits_bytes) // (max(seq_len_sum, 1) * 4)) block_n = 32 select_warps = 8 - select_block = _next_power_of_2(max(seq_len_sum, topk)) + 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( 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 index c984ff85a..fd9dd723c 100644 --- a/tokenspeed-kernel-amd/test/ops/test_attention_dsa_gluon_gfx950.py +++ b/tokenspeed-kernel-amd/test/ops/test_attention_dsa_gluon_gfx950.py @@ -514,6 +514,28 @@ def _assert_topk_matches( 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, @@ -566,6 +588,57 @@ def test_dsa_decode_topk_fp8_glm52_cases(case: _TopKDecodeCase) -> None: _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,