From 530965688d50b552ea9351dc3608099441b4a8dd Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Sat, 14 Nov 2026 00:40:10 +0000 Subject: [PATCH] fix(attention): bound multi-row triton decode tiles by the LDS size on ROCm --- python/freetoken/kernel/triton/attention.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index f6a84692e..385ee6cdd 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -92,6 +92,11 @@ def _optin_smem_bytes(device_index: int) -> int: return int(getattr(props, "shared_memory_per_block_optin", 0)) +@functools.lru_cache(maxsize=None) +def _rocm_lds_bytes(device_index: int) -> int: + return int(torch.cuda.get_device_properties(device_index).shared_memory_per_block) + + @functools.lru_cache(maxsize=None) def _num_sms(device_index: int) -> int: return int(torch.cuda.get_device_properties(device_index).multi_processor_count) @@ -612,8 +617,10 @@ def decode_paged_attention( if kv_format.quantized: block_n, num_stages, num_warps = _quantized_decode_launch(block_d, block_n, num_stages) if rows > 1: + # NOTE: torch reads 0 opt-in shared memory on ROCm, which the launch rule takes as no limit + smem_budget = _rocm_lds_bytes(q.device.index) if is_rocm() else smem block_n, num_stages, num_warps = _multi_row_decode_launch( - rows * block_h, block_d, kv_format.quantized, elem, k_cache.element_size(), smem + rows * block_h, block_d, kv_format.quantized, elem, k_cache.element_size(), smem_budget ) k_scale_arg, v_scale_arg, stride_kss, stride_vss = kv_scale_args(kv_format, k_cache, v_cache, k_scale, v_scale, q)