From 2618f682e6c4763b916f07f7b08811b7f120006e Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Fri, 9 Oct 2026 16:04:11 -0500 Subject: [PATCH 01/12] fix: warmup kernels Signed-off-by: simondanielsson --- .../layers/attention/mla_attention.py | 6 +- .../models/glm5next/amd/ops/kpool_compress.py | 168 +++++++++++++++--- vllm/models/glm5next/amd/sparse_indexer.py | 15 ++ vllm/models/glm5next/common/kda.py | 6 + .../flash_linear_attention/ops/l2norm.py | 54 ++++-- 5 files changed, 209 insertions(+), 40 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 2a3048096c98..09664afe1fea 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -662,6 +662,7 @@ def __init__( "FLASHINFER_MLA_SPARSE", "FLASHINFER_MLA_SPARSE_SM120", "DEEPSEEK_V32_INDEXER", + "ROCM_AITER_MLA_SPARSE", ): from vllm.v1.attention.backends.mla.compressor_utils import ( _COMPRESSED_SLOT_MAPPING_KERNEL, @@ -675,7 +676,10 @@ def __init__( _PREPARE_UNIFORM_DECODE_KERNEL.register_warmup() _BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup() - if backend_name != "DEEPSEEK_V32_INDEXER": + if backend_name not in ( + "DEEPSEEK_V32_INDEXER", + "ROCM_AITER_MLA_SPARSE", + ): from vllm.v1.attention.backends.mla.sparse_swa import ( _COMPUTE_PREFILL_METADATA_KERNEL, ) diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index 766950f3321c..18fff126132f 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -12,8 +12,16 @@ from __future__ import annotations +from typing import Any + import torch +from vllm.model_executor.warmup.jit_warmup import WarmupChoices +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + DispatchSpec, + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -191,6 +199,78 @@ def _kpool_softmax_rotate_write_cache_kernel( tl.store(compressed_scale_ptr + row, scale) +def _kpool_compress_warmup_inputs( + *, k_cache, pool_size: int, head_dim: int, round_scale: bool +) -> dict[str, Any]: + kv_cache = k_cache.kv_cache + page_size = kv_cache.shape[1] + buf_numel_per_page = kv_cache.stride(0) + loc_aligned: Any = WarmupChoices(True, False) + return dict( + buf_fp8_ptr=TritonWarmupTensor(FP8_DTYPE), + buf_fp32_ptr=TritonWarmupTensor(torch.float32), + slot_k_ptr=TritonWarmupTensor(torch.bfloat16, shape=(1, pool_size, head_dim)), + slot_score_ptr=TritonWarmupTensor( + torch.bfloat16, shape=(1, pool_size, head_dim) + ), + ape_ptr=TritonWarmupTensor(torch.float32, shape=(pool_size, head_dim)), + loc_ptr=TritonWarmupTensor(torch.int64, aligned=loc_aligned), + write_mask_ptr=TritonWarmupTensor(torch.bool), + compressed_k_ptr=TritonWarmupTensor(FP8_DTYPE), + compressed_scale_ptr=TritonWarmupTensor(torch.float32), + page_size=page_size, + buf_numel_per_page=buf_numel_per_page, + head_dim=head_dim, + round_scale=round_scale, + has_write_mask=True, + return_compressed=False, + write_cache=True, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_kpool_softmax_rotate_write_cache_kernel, + warmup_inputs=_kpool_compress_warmup_inputs, +) +def kpool_compress_dispatch( + buf_fp8_ptr, + buf_fp32_ptr, + slot_k_ptr, + slot_score_ptr, + ape_ptr, + loc_ptr, + write_mask_ptr, + compressed_k_ptr, + compressed_scale_ptr, + page_size: int, + buf_numel_per_page: int, + head_dim: int, + round_scale: bool, + has_write_mask: bool, + return_compressed: bool, + write_cache: bool, +) -> DispatchSpec: + return (slot_k_ptr.shape[0],), dict( + slot_k_stride_0=slot_k_ptr.stride(0), + slot_k_stride_1=slot_k_ptr.stride(1), + slot_score_stride_0=slot_score_ptr.stride(0), + slot_score_stride_1=slot_score_ptr.stride(1), + ape_stride_0=ape_ptr.stride(0), + PAGE_SIZE=page_size, + BUF_NUMEL_PER_PAGE=buf_numel_per_page, + POOL_SIZE=slot_k_ptr.shape[1], + HEAD_DIM=head_dim, + S_OFFSET_NBYTES_IN_PAGE=page_size * head_dim, + FP8_MAX=FP8_MAX, + PRESHUFFLE=page_size > 1, + ROUND_SCALE=round_scale, + HAS_WRITE_MASK=has_write_mask, + RETURN_COMPRESSED=return_compressed, + WRITE_CACHE=write_cache, + BLOCK_D=triton.next_power_of_2(head_dim), + ) + + def kpool_compress_and_write_cache( kv_cache: torch.Tensor, slot_k: torch.Tensor, @@ -261,7 +341,6 @@ def kpool_compress_and_write_cache( buf_fp32 = buf.view(torch.float32) # bytes per page (last dim of kv_cache) viewed as uint8 buf_numel_per_page = buf.stride(0) - s_offset_nbytes_in_page = page_size * head_dim if return_compressed: compressed_k = torch.empty( @@ -279,7 +358,7 @@ def kpool_compress_and_write_cache( if page_size > 1: assert page_size % 16 == 0, "ROCm preshuffle requires 16-token tiles" - _kpool_softmax_rotate_write_cache_kernel[(slot_k.shape[0],)]( + kpool_compress_dispatch( buf_fp8, buf_fp32, slot_k, @@ -289,23 +368,13 @@ def kpool_compress_and_write_cache( write_mask, compressed_k, compressed_scale, - slot_k.stride(0), - slot_k.stride(1), - slot_score.stride(0), - slot_score.stride(1), - ape.stride(0), - PAGE_SIZE=page_size, - BUF_NUMEL_PER_PAGE=buf_numel_per_page, - POOL_SIZE=slot_k.shape[1], - HEAD_DIM=head_dim, - S_OFFSET_NBYTES_IN_PAGE=s_offset_nbytes_in_page, - FP8_MAX=FP8_MAX, - PRESHUFFLE=page_size > 1, - ROUND_SCALE=round_scale, - HAS_WRITE_MASK=has_write_mask, - RETURN_COMPRESSED=return_compressed, - WRITE_CACHE=write_cache, - BLOCK_D=triton.next_power_of_2(head_dim), + page_size, + buf_numel_per_page, + head_dim, + round_scale, + has_write_mask, + return_compressed, + write_cache, ) if return_compressed: @@ -362,6 +431,54 @@ def _kpool_tail_seed_kernel( ) +def _kpool_tail_seed_warmup_inputs( + *, tail_cache, kpool: int, head_dim: int +) -> dict[str, Any]: + tail_kv_cache = tail_cache.kv_cache + tail_block_elems = tail_kv_cache.stride(0) + kpool_head = tail_kv_cache.stride(1) + ring = tail_kv_cache.shape[2] + tslot_aligned: Any = WarmupChoices(True, False) + return dict( + key_ptr=TritonWarmupTensor(torch.bfloat16), + score_ptr=TritonWarmupTensor(torch.bfloat16), + tslot_ptr=TritonWarmupTensor(torch.int64, aligned=tslot_aligned), + tail_ptr=TritonWarmupTensor(torch.bfloat16), + n_tokens=WarmupChoices(1, 16, 2), + tail_block_elems=tail_block_elems, + kpool_head=kpool_head, + head_dim=head_dim, + kpool=kpool, + ring=ring, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=_kpool_tail_seed_kernel, + warmup_inputs=_kpool_tail_seed_warmup_inputs, +) +def kpool_tail_seed_dispatch( + key_ptr, + score_ptr, + tslot_ptr, + tail_ptr, + n_tokens: int, + tail_block_elems: int, + kpool_head: int, + head_dim: int, + kpool: int, + ring: int, +) -> DispatchSpec: + return (n_tokens,), dict( + TAIL_BLOCK_ELEMS=tail_block_elems, + KPOOL_HEAD=kpool_head, + HEAD_DIM=head_dim, + KPOOL=kpool, + RING=ring, + BLOCK_D=triton.next_power_of_2(head_dim), + ) + + def kpool_seed_tail_cache( tail_kv_cache: torch.Tensor, key: torch.Tensor, @@ -376,18 +493,17 @@ def kpool_seed_tail_cache( n = tslot.shape[0] if n == 0: return - _kpool_tail_seed_kernel[(n,)]( + kpool_tail_seed_dispatch( key, gate_score, tslot, tail_kv_cache, n, - TAIL_BLOCK_ELEMS=tail_kv_cache.stride(0), - KPOOL_HEAD=tail_kv_cache.stride(1), - HEAD_DIM=head_dim, - KPOOL=kpool, - RING=tail_kv_cache.shape[2], - BLOCK_D=triton.next_power_of_2(head_dim), + tail_kv_cache.stride(0), + tail_kv_cache.stride(1), + head_dim, + kpool, + tail_kv_cache.shape[2], ) diff --git a/vllm/models/glm5next/amd/sparse_indexer.py b/vllm/models/glm5next/amd/sparse_indexer.py index 13e67de04bf7..a88b6b0b7a2f 100644 --- a/vllm/models/glm5next/amd/sparse_indexer.py +++ b/vllm/models/glm5next/amd/sparse_indexer.py @@ -737,6 +737,21 @@ def __init__( raise NotImplementedError( "SparseAttnIndexerKpool on ROCm does not support DCP." ) + hf_config = cfg.model_config.hf_text_config if cfg is not None else None + index_kpool = getattr(hf_config, "index_kpool", None) or 1 + if index_kpool > 1: + kpool_ops.kpool_compress_dispatch.register_warmup( + k_cache=k_cache, + pool_size=index_kpool, + head_dim=head_dim, + round_scale=scale_fmt is not None, + ) + if tail_cache is not None: + kpool_ops.kpool_tail_seed_dispatch.register_warmup( + tail_cache=tail_cache, + kpool=index_kpool, + head_dim=head_dim, + ) def forward_hip( self, diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 141a8cc7cb82..383dca0141ac 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -45,6 +45,7 @@ ) from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated +from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd_block from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.worker.workspace import current_workspace_manager @@ -347,6 +348,11 @@ def _a_log_weight_loader(param, loaded_weight): vllm_config.model_config.dtype, self.kda_lower_bound, ) + if self.kda_prefill_backend != "flashkda": + l2norm_fwd_block.register_warmup( + dtype=vllm_config.model_config.dtype, + dim=self.head_dim, + ) self._flashkda_buffer_specs: ( tuple[tuple[tuple[int, ...], torch.dtype], ...] | None ) = None diff --git a/vllm/third_party/flash_linear_attention/ops/l2norm.py b/vllm/third_party/flash_linear_attention/ops/l2norm.py index 2eb137a242fb..c2f5a2ea2882 100644 --- a/vllm/third_party/flash_linear_attention/ops/l2norm.py +++ b/vllm/third_party/flash_linear_attention/ops/l2norm.py @@ -8,15 +8,24 @@ # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang import os +from typing import Any import torch +from vllm.model_executor.warmup.jit_warmup import WarmupChoices +from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + DispatchSpec, + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, +) from vllm.triton_utils import tl, triton BT_LIST = [8, 16, 32, 64, 128] USE_DEFAULT_FLA_NORM = int(os.getenv("USE_DEFAULT_FLA_NORM", "0")) +_L2NORM_MBLOCK = 32 + @triton.autotune( configs=[ @@ -92,6 +101,36 @@ def l2norm_fwd_kernel2( tl.store(Y + (rindex + N * row_idx), xs * rsqrt, mask) +def _l2norm_block_dim(element_size: int, dim: int) -> int: + return min(65536 // element_size, triton.next_power_of_2(dim)) + + +def _l2norm_fwd_block_warmup_inputs(*, dtype: torch.dtype, dim: int) -> dict[str, Any]: + block_dim = _l2norm_block_dim(dtype.itemsize, dim) + num_rows: Any = WarmupChoices(1, 16, 2) + x_aligned: Any = WarmupChoices(True, False) + return dict( + x=TritonWarmupTensor(dtype, aligned=x_aligned, shape=(num_rows, dim)), + y=TritonWarmupTensor(dtype), + eps=1e-6, + block_dim=block_dim, + ) + + +@triton_kernel_dispatcher_with_warmup( + kernel=l2norm_fwd_kernel2, + warmup_inputs=_l2norm_fwd_block_warmup_inputs, +) +def l2norm_fwd_block(x, y, eps: float, block_dim: int) -> DispatchSpec: + num_rows, dim = x.shape[0], x.shape[-1] + return (triton.cdiv(num_rows, _L2NORM_MBLOCK),), dict( + M=num_rows, + N=dim, + BD=block_dim, + MBLOCK=_L2NORM_MBLOCK, + ) + + def l2norm_fwd( x: torch.Tensor, eps: float = 1e-6, output_dtype: torch.dtype | None = None ): @@ -106,23 +145,12 @@ def l2norm_fwd( T, D = x.shape[0], x.shape[-1] # rstd = torch.empty((T,), dtype=torch.float32, device=x.device) # Less than 64KB per feature: enqueue fused kernel - MAX_FUSED_SIZE = 65536 // x.element_size() - BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) + BD = _l2norm_block_dim(x.element_size(), D) if D > BD: raise RuntimeError("This layer doesn't support feature dim >= 64KB.") if not USE_DEFAULT_FLA_NORM: - MBLOCK = 32 - # M, N = x.shape - l2norm_fwd_kernel2[(triton.cdiv(T, MBLOCK),)]( - x, - y, - eps, - T, - D, - BD, - MBLOCK, - ) + l2norm_fwd_block(x, y, eps, BD) else: if D <= 512: NB = triton.cdiv(T, 2048) From 20b45154273c648448490ccdfc886a5928eef6da Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 10:52:31 -0500 Subject: [PATCH 02/12] fix: fix speccialization on 2gb+ buffers Signed-off-by: simondanielsson --- vllm/model_executor/warmup/jit_warmup_triton_helper.py | 10 +++++++++- vllm/models/glm5next/amd/ops/kpool_compress.py | 8 +++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index 73db12318b1a..095f0df6e9e4 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -174,6 +174,11 @@ class TritonWarmupTensor: ``strides=None`` represents compact row-major storage. Pass explicit strides whenever the runtime tensor can be padded, transposed, or otherwise strided. + + ``storage_bytes`` must be set from ``untyped_storage().size()`` of the + runtime tensor whenever that storage can exceed 2 GiB. The ROCm backend + specializes pointers on it (``tt.pointer_range``), so a KV-cache view left + at the 0 default value warms a variant the live launch never uses. """ dtype: Any @@ -181,12 +186,13 @@ class TritonWarmupTensor: shape: tuple[int, ...] = (1,) strides: tuple[int, ...] | None = None init: Any = 0 + storage_bytes: int = 0 def data_ptr(self) -> int: return 0 if self.aligned else 1 def ptr_range(self) -> int: - return 0 + return self.storage_bytes @property def device(self) -> Any: @@ -497,6 +503,7 @@ def _triton_metadata_arg(value: Any) -> Any: aligned=getattr(value, "_vllm_warmup_aligned", True), shape=tuple(value.shape), strides=tuple(value.stride()), + storage_bytes=getattr(value, "_vllm_warmup_storage_bytes", 0), ) @@ -536,6 +543,7 @@ def materialize(value: Any) -> Any: tensor.fill_(value.init) else: tensor._vllm_warmup_aligned = value.aligned + tensor._vllm_warmup_storage_bytes = value.storage_bytes with suppress(TypeError): cache[value] = tensor return tensor diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index 18fff126132f..86974af9417a 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -205,10 +205,11 @@ def _kpool_compress_warmup_inputs( kv_cache = k_cache.kv_cache page_size = kv_cache.shape[1] buf_numel_per_page = kv_cache.stride(0) + buf_storage_bytes = kv_cache.untyped_storage().size() loc_aligned: Any = WarmupChoices(True, False) return dict( - buf_fp8_ptr=TritonWarmupTensor(FP8_DTYPE), - buf_fp32_ptr=TritonWarmupTensor(torch.float32), + buf_fp8_ptr=TritonWarmupTensor(FP8_DTYPE, storage_bytes=buf_storage_bytes), + buf_fp32_ptr=TritonWarmupTensor(torch.float32, storage_bytes=buf_storage_bytes), slot_k_ptr=TritonWarmupTensor(torch.bfloat16, shape=(1, pool_size, head_dim)), slot_score_ptr=TritonWarmupTensor( torch.bfloat16, shape=(1, pool_size, head_dim) @@ -438,12 +439,13 @@ def _kpool_tail_seed_warmup_inputs( tail_block_elems = tail_kv_cache.stride(0) kpool_head = tail_kv_cache.stride(1) ring = tail_kv_cache.shape[2] + tail_storage_bytes = tail_kv_cache.untyped_storage().size() tslot_aligned: Any = WarmupChoices(True, False) return dict( key_ptr=TritonWarmupTensor(torch.bfloat16), score_ptr=TritonWarmupTensor(torch.bfloat16), tslot_ptr=TritonWarmupTensor(torch.int64, aligned=tslot_aligned), - tail_ptr=TritonWarmupTensor(torch.bfloat16), + tail_ptr=TritonWarmupTensor(torch.bfloat16, storage_bytes=tail_storage_bytes), n_tokens=WarmupChoices(1, 16, 2), tail_block_elems=tail_block_elems, kpool_head=kpool_head, From 5f772df259344dcd7bd657b1bdd0a482b024575a Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 10:55:37 -0500 Subject: [PATCH 03/12] fix: add util Signed-off-by: simondanielsson --- .../warmup/jit_warmup_triton_helper.py | 20 +++++++++++++++---- .../models/glm5next/amd/ops/kpool_compress.py | 5 +++-- 2 files changed, 19 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index 095f0df6e9e4..d76a93def37c 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -168,6 +168,18 @@ def triton_scalar_specialization_rep(value: int) -> int: return divisible_rep if value % 16 == 0 else generic_rep +def triton_ptr_range(tensor: Any) -> int: + """Pointer range Triton's ROCm backend would read off ``tensor``. + + Mirrors ``HIPBackend.is_within_2gb``: an explicit ``ptr_range`` wins over + the storage size, so a view that narrows its own range stays narrow here. + """ + ptr_range = getattr(tensor, "ptr_range", None) + if ptr_range is not None: + return ptr_range() + return tensor.untyped_storage().size() + + @dataclass(frozen=True) class TritonWarmupTensor: """Compile-only tensor metadata used by Triton warmup. @@ -175,10 +187,10 @@ class TritonWarmupTensor: ``strides=None`` represents compact row-major storage. Pass explicit strides whenever the runtime tensor can be padded, transposed, or otherwise strided. - ``storage_bytes`` must be set from ``untyped_storage().size()`` of the - runtime tensor whenever that storage can exceed 2 GiB. The ROCm backend - specializes pointers on it (``tt.pointer_range``), so a KV-cache view left - at the 0 default value warms a variant the live launch never uses. + ``storage_bytes`` must be set via ``triton_ptr_range`` whenever the runtime + tensor's storage can exceed 2 GiB. The ROCm backend specializes pointers on + it (``tt.pointer_range``), so a KV-cache view left at the 0 default value + warms a variant the live launch never uses. """ dtype: Any diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index 86974af9417a..b9c7208b0dfb 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -21,6 +21,7 @@ DispatchSpec, TritonWarmupTensor, triton_kernel_dispatcher_with_warmup, + triton_ptr_range, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -205,7 +206,7 @@ def _kpool_compress_warmup_inputs( kv_cache = k_cache.kv_cache page_size = kv_cache.shape[1] buf_numel_per_page = kv_cache.stride(0) - buf_storage_bytes = kv_cache.untyped_storage().size() + buf_storage_bytes = triton_ptr_range(kv_cache) loc_aligned: Any = WarmupChoices(True, False) return dict( buf_fp8_ptr=TritonWarmupTensor(FP8_DTYPE, storage_bytes=buf_storage_bytes), @@ -439,7 +440,7 @@ def _kpool_tail_seed_warmup_inputs( tail_block_elems = tail_kv_cache.stride(0) kpool_head = tail_kv_cache.stride(1) ring = tail_kv_cache.shape[2] - tail_storage_bytes = tail_kv_cache.untyped_storage().size() + tail_storage_bytes = triton_ptr_range(tail_kv_cache) tslot_aligned: Any = WarmupChoices(True, False) return dict( key_ptr=TritonWarmupTensor(torch.bfloat16), From 6839d615243e58f09ae5f493ca35a8cffc4230c1 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 11:13:38 -0500 Subject: [PATCH 04/12] fix: improve docstring Signed-off-by: simondanielsson --- vllm/model_executor/warmup/jit_warmup_triton_helper.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index d76a93def37c..ed9fc7d06a2d 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -187,10 +187,11 @@ class TritonWarmupTensor: ``strides=None`` represents compact row-major storage. Pass explicit strides whenever the runtime tensor can be padded, transposed, or otherwise strided. - ``storage_bytes`` must be set via ``triton_ptr_range`` whenever the runtime - tensor's storage can exceed 2 GiB. The ROCm backend specializes pointers on - it (``tt.pointer_range``), so a KV-cache view left at the 0 default value - warms a variant the live launch never uses. + On platforms where triton specializes on buffer sizes (ROCm): + ``storage_bytes`` should be set via ``triton_ptr_range`` whenever the runtime + tensor's storage can exceed 2 GiB. The ROCm backend specializes pointers on + it (``tt.pointer_range``) to enable lowering to buffer load/store intrinsics. + Otherwise warmup may compile a variant that the live launch never uses. """ dtype: Any From 14b373e09b8dd514c1b1c5f2bedc5182fa333b21 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 11:15:21 -0500 Subject: [PATCH 05/12] chore: simplify another docstring Signed-off-by: simondanielsson --- vllm/model_executor/warmup/jit_warmup_triton_helper.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index ed9fc7d06a2d..8d8b81bda434 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -169,10 +169,9 @@ def triton_scalar_specialization_rep(value: int) -> int: def triton_ptr_range(tensor: Any) -> int: - """Pointer range Triton's ROCm backend would read off ``tensor``. + """Pointer range Triton's backend would read off ``tensor``. - Mirrors ``HIPBackend.is_within_2gb``: an explicit ``ptr_range`` wins over - the storage size, so a view that narrows its own range stays narrow here. + Mirrors ``HIPBackend.is_within_2gb``. """ ptr_range = getattr(tensor, "ptr_range", None) if ptr_range is not None: From 67cb69504043157a1e95bae31c63a8a726183dee Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 11:26:01 -0500 Subject: [PATCH 06/12] feat: add regression test Signed-off-by: simondanielsson --- .../test_jit_warmup_triton_launcher.py | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index 0e3ff2b0c88e..6ef4cd36bd78 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/tests/model_executor/test_jit_warmup_triton_launcher.py @@ -22,6 +22,7 @@ TritonWarmupTensor, VllmTritonJitKernel, triton_kernel_dispatcher_with_warmup, + triton_ptr_range, triton_warmup_inputs, ) from vllm.platforms import current_platform @@ -479,6 +480,37 @@ def test_triton_warmup_inputs_expands_explicit_pointer_dtypes() -> None: ) +@pytest.mark.parametrize("live_ptr_range", [None, 1 << 32]) +def test_warmup_pointer_specializes_like_the_live_tensor( + live_ptr_range: int | None, +) -> None: + """A stand-in must land on the same side of ROCm's 2 GiB buffer-ops window. + + KV-cache views sit in an allocation far past it, so a stand-in that keeps + the optimistic default range warms a ``tt.pointer_range`` specialization no + launch ever asks for, and the kernel JITs on its first real batch. + """ + backend = pytest.importorskip("triton.backends.amd.compiler").HIPBackend + if not hasattr(backend, "is_within_2gb"): + # if it ever changes in the future + assert not hasattr(backend, "get_tensor_specialization"), ( + "Triton still specializes tensor arguments but dropped " + "is_within_2gb; re-derive TritonWarmupTensor.storage_bytes " + "against whatever now decides tt.pointer_range" + ) + pytest.skip("Triton no longer specializes tensor arguments") + is_within_2gb = backend.is_within_2gb + + live = torch.empty(64, dtype=torch.uint8) + if live_ptr_range is not None: + live.ptr_range = lambda: live_ptr_range + + stub = TritonWarmupTensor(torch.bfloat16, storage_bytes=triton_ptr_range(live)) + + assert is_within_2gb(stub) == is_within_2gb(live) + assert is_within_2gb(stub) is (live_ptr_range is None) + + def test_triton_key_derivation_applies_wrappers_and_runtime_options( monkeypatch: pytest.MonkeyPatch, ) -> None: From bc082823a8f0fa3addb15cb7a0d95971d35e6a91 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Tue, 20 Oct 2026 11:57:32 -0500 Subject: [PATCH 07/12] fix: set storage_bytes also onthe compressed k and sclae ptrs Signed-off-by: simondanielsson --- vllm/models/glm5next/amd/ops/kpool_compress.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index b9c7208b0dfb..debb8a73b9a7 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -208,9 +208,12 @@ def _kpool_compress_warmup_inputs( buf_numel_per_page = kv_cache.stride(0) buf_storage_bytes = triton_ptr_range(kv_cache) loc_aligned: Any = WarmupChoices(True, False) + # With return_compressed=False the compressed outputs alias the cache views. + buf_fp8 = TritonWarmupTensor(FP8_DTYPE, storage_bytes=buf_storage_bytes) + buf_fp32 = TritonWarmupTensor(torch.float32, storage_bytes=buf_storage_bytes) return dict( - buf_fp8_ptr=TritonWarmupTensor(FP8_DTYPE, storage_bytes=buf_storage_bytes), - buf_fp32_ptr=TritonWarmupTensor(torch.float32, storage_bytes=buf_storage_bytes), + buf_fp8_ptr=buf_fp8, + buf_fp32_ptr=buf_fp32, slot_k_ptr=TritonWarmupTensor(torch.bfloat16, shape=(1, pool_size, head_dim)), slot_score_ptr=TritonWarmupTensor( torch.bfloat16, shape=(1, pool_size, head_dim) @@ -218,8 +221,8 @@ def _kpool_compress_warmup_inputs( ape_ptr=TritonWarmupTensor(torch.float32, shape=(pool_size, head_dim)), loc_ptr=TritonWarmupTensor(torch.int64, aligned=loc_aligned), write_mask_ptr=TritonWarmupTensor(torch.bool), - compressed_k_ptr=TritonWarmupTensor(FP8_DTYPE), - compressed_scale_ptr=TritonWarmupTensor(torch.float32), + compressed_k_ptr=buf_fp8, + compressed_scale_ptr=buf_fp32, page_size=page_size, buf_numel_per_page=buf_numel_per_page, head_dim=head_dim, From 485523622f051bdb6cc1ad540a9dd3160ce8c05e Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 28 Oct 2026 08:59:05 -0500 Subject: [PATCH 08/12] test: add test that ensures the compile call propagates the ptr range bytes Signed-off-by: simondanielsson --- .../test_jit_warmup_triton_launcher.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index e59644105fc6..99d31d041844 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/tests/model_executor/test_jit_warmup_triton_launcher.py @@ -511,6 +511,49 @@ def test_warmup_pointer_specializes_like_the_live_tensor( assert is_within_2gb(stub) is (live_ptr_range is None) +def test_warmup_pointer_range_survives_the_fake_tensor_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A dispatcher that needs real tensors must not drop ``ptr_range_bytes``. + + A dispatcher touching tensor APIs re-runs its case through fake tensors, + so the declared range has to round-trip the stand-in as an attribute and + come back out of both conversion points, key derivation and the warmup + launch. Dropping it there silently restores the optimistic <2 GiB + specialization and the kernel JITs on its first real batch. + """ + kernel = _FakeTritonKernel() + over_2gb = 1 << 32 + + def warmup_inputs() -> dict[str, Any]: + return dict( + first=TritonWarmupTensor( + torch.bfloat16, shape=(2, 3), ptr_range_bytes=over_2gb + ), + second=1, + ) + + @triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs) + def dispatch(first: torch.Tensor, second: int) -> LaunchSpec: + assert isinstance(first, torch.Tensor) + return (first.shape[0],), dict(CONST=first[0].numel()) + + derived: list[dict[str, Any]] = [] + + def derive(kernel: Any, kwargs: dict[str, Any]) -> set[TritonJitKey]: + derived.append(kwargs) + return {TritonJitKey(id(kernel), "fake", 0, kwargs["CONST"])} + + _patch_key_deriver(monkeypatch, derive) + + keys = dispatch.get_warmup_keys() + kernel.warmup_calls.clear() + dispatch.compile(keys[0]) + + assert [kwargs["first"].ptr_range_bytes for kwargs in derived] == [over_2gb] + assert [call["first"].ptr_range_bytes for call in kernel.warmup_calls] == [over_2gb] + + def test_triton_key_derivation_applies_wrappers_and_runtime_options( monkeypatch: pytest.MonkeyPatch, ) -> None: From 29b1816aebf6526bae67ce38e67b0c6a99c80aaf Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 28 Oct 2026 11:46:40 -0500 Subject: [PATCH 09/12] fix: only warmup l2norm on triton backend Signed-off-by: simondanielsson --- vllm/models/glm5next/common/kda.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/models/glm5next/common/kda.py b/vllm/models/glm5next/common/kda.py index 28c816c3986b..3d2febe95761 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -359,7 +359,7 @@ def _a_log_weight_loader(param, loaded_weight): vllm_config.model_config.dtype, self.kda_lower_bound, ) - if self.kda_prefill_backend != "flashkda": + if self.kda_prefill_backend == "triton": l2norm_fwd_block.register_warmup( dtype=vllm_config.model_config.dtype, dim=self.head_dim, From 51d7cd814fe6b4ae899ce69df918659e6bcac6e0 Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Wed, 28 Oct 2026 11:50:42 -0500 Subject: [PATCH 10/12] fix: remove change to mla_attention Signed-off-by: simondanielsson --- vllm/model_executor/layers/attention/mla_attention.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 2adeef127434..f64f989434c9 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -681,10 +681,7 @@ def __init__( _PREPARE_UNIFORM_DECODE_KERNEL.register_warmup() _BUILD_PREFILL_CHUNK_METADATA_KERNEL.register_warmup() - if backend_name not in ( - "DEEPSEEK_V32_INDEXER", - "ROCM_AITER_MLA_SPARSE", - ): + if backend_name != "DEEPSEEK_V32_INDEXER": from vllm.v1.attention.backends.mla.sparse_swa import ( _COMPUTE_PREFILL_METADATA_KERNEL, ) From 00c873fa476673a500e47b77a1091b4e0d9e08dd Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Thu, 29 Oct 2026 14:56:54 -0500 Subject: [PATCH 11/12] fix: make sure tests works even if no modelconfig is none Signed-off-by: simondanielsson --- vllm/models/glm5next/amd/sparse_indexer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/models/glm5next/amd/sparse_indexer.py b/vllm/models/glm5next/amd/sparse_indexer.py index c5a48baf7552..9b9920852f9f 100644 --- a/vllm/models/glm5next/amd/sparse_indexer.py +++ b/vllm/models/glm5next/amd/sparse_indexer.py @@ -754,7 +754,8 @@ def __init__( raise NotImplementedError( "SparseAttnIndexerKpool on ROCm does not support DCP." ) - hf_config = cfg.model_config.hf_text_config if cfg is not None else None + model_config = cfg.model_config if cfg is not None else None + hf_config = model_config.hf_text_config if model_config is not None else None index_kpool = getattr(hf_config, "index_kpool", None) or 1 if index_kpool > 1: kpool_ops.kpool_compress_dispatch.register_warmup( From ec18e933a78d4e13debcc78ac1ac82f5e9239b7a Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Fri, 30 Oct 2026 14:02:15 -0500 Subject: [PATCH 12/12] fix: don't specialize redundantly on alignment for sccalar load pointers Signed-off-by: simondanielsson --- vllm/models/glm5next/amd/ops/kpool_compress.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/vllm/models/glm5next/amd/ops/kpool_compress.py b/vllm/models/glm5next/amd/ops/kpool_compress.py index cbeac1695508..fa38e8a86a28 100644 --- a/vllm/models/glm5next/amd/ops/kpool_compress.py +++ b/vllm/models/glm5next/amd/ops/kpool_compress.py @@ -16,7 +16,6 @@ import torch -from vllm.model_executor.warmup.jit_warmup import WarmupChoices from vllm.model_executor.warmup.jit_warmup_triton_helper import ( DispatchSpec, TritonWarmupTensor, @@ -77,7 +76,7 @@ def _hadamard128(x): # Fused pool compression and cache write. -@triton.jit +@triton.jit(do_not_specialize_on_alignment=["loc_ptr"]) def _kpool_softmax_rotate_write_cache_kernel( buf_fp8_ptr, buf_fp32_ptr, @@ -207,7 +206,6 @@ def _kpool_compress_warmup_inputs( page_size = kv_cache.shape[1] buf_numel_per_page = kv_cache.stride(0) buf_ptr_range_bytes = triton_ptr_range(kv_cache) - loc_aligned: Any = WarmupChoices(True, False) # With return_compressed=False the compressed outputs alias the cache views. buf_fp8 = TritonWarmupTensor(FP8_DTYPE, ptr_range_bytes=buf_ptr_range_bytes) buf_fp32 = TritonWarmupTensor(torch.float32, ptr_range_bytes=buf_ptr_range_bytes) @@ -219,7 +217,7 @@ def _kpool_compress_warmup_inputs( torch.bfloat16, shape=(1, pool_size, head_dim) ), ape_ptr=TritonWarmupTensor(torch.float32, shape=(pool_size, head_dim)), - loc_ptr=TritonWarmupTensor(torch.int64, aligned=loc_aligned), + loc_ptr=TritonWarmupTensor(torch.int64), write_mask_ptr=TritonWarmupTensor(torch.bool), compressed_k_ptr=buf_fp8, compressed_scale_ptr=buf_fp32, @@ -390,7 +388,9 @@ def kpool_compress_and_write_cache( # Seed each request's incomplete pool into its paged tail during prefill. -@triton.jit +@triton.jit( + do_not_specialize=["n_tokens"], do_not_specialize_on_alignment=["tslot_ptr"] +) def _kpool_tail_seed_kernel( key_ptr, score_ptr, @@ -444,15 +444,14 @@ def _kpool_tail_seed_warmup_inputs( kpool_head = tail_kv_cache.stride(1) ring = tail_kv_cache.shape[2] tail_ptr_range_bytes = triton_ptr_range(tail_kv_cache) - tslot_aligned: Any = WarmupChoices(True, False) return dict( key_ptr=TritonWarmupTensor(torch.bfloat16), score_ptr=TritonWarmupTensor(torch.bfloat16), - tslot_ptr=TritonWarmupTensor(torch.int64, aligned=tslot_aligned), + tslot_ptr=TritonWarmupTensor(torch.int64), tail_ptr=TritonWarmupTensor( torch.bfloat16, ptr_range_bytes=tail_ptr_range_bytes ), - n_tokens=WarmupChoices(1, 16, 2), + n_tokens=1, tail_block_elems=tail_block_elems, kpool_head=kpool_head, head_dim=head_dim,