diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index 0e3ff2b0c88e..99d31d041844 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,80 @@ 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.ptr_range_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, ptr_range_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_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: diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 5eef73e49c4a..f64f989434c9 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -667,6 +667,7 @@ def __init__( "FLASHINFER_MLA_SPARSE_SM120", "FLASH_ATTN_MLA_SPARSE_FA4", "DEEPSEEK_V32_INDEXER", + "ROCM_AITER_MLA_SPARSE", ): from vllm.v1.attention.backends.mla.compressor_utils import ( _COMPRESSED_SLOT_MAPPING_KERNEL, diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index 1fdaa20e8706..ea4ffacbd280 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -168,6 +168,17 @@ 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 backend would read off ``tensor``. + + Mirrors ``HIPBackend.is_within_2gb``. + """ + 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. @@ -176,12 +187,13 @@ class TritonWarmupTensor: whenever the runtime tensor can be padded, transposed, or otherwise strided. ``ptr_range_bytes`` is the byte size of the storage behind the runtime - tensor (its ``untyped_storage()``, not the view). With buffer ops enabled, - Triton's ROCm backend gives pointers into storage of at most 2 GB a - ``tt.pointer_range = 32`` specialization (32-bit buffer offsets), which is - part of the compile key. The default 0 assumes such a small storage; set it - above 2 GB for tensors that view a larger pool, such as a KV cache layer, - or warmup compiles a variant the runtime never uses. CUDA ignores it. + tensor (its ``untyped_storage()``, not the view), most easily derived with + ``triton_ptr_range``. With buffer ops enabled, Triton's ROCm backend gives + pointers into storage of at most 2 GB a ``tt.pointer_range = 32`` + specialization (32-bit buffer offsets), which is part of the compile key. + The default 0 assumes such a small storage; set it above 2 GB for tensors + that view a larger pool, such as a KV cache layer, or warmup compiles a + variant the runtime never uses. CUDA ignores it. """ dtype: Any @@ -507,6 +519,7 @@ def _triton_metadata_arg(value: Any) -> Any: aligned=getattr(value, "_vllm_warmup_aligned", True), shape=tuple(value.shape), strides=tuple(value.stride()), + ptr_range_bytes=getattr(value, "_vllm_warmup_ptr_range_bytes", 0), ) @@ -546,6 +559,7 @@ def materialize(value: Any) -> Any: tensor.fill_(value.init) else: tensor._vllm_warmup_aligned = value.aligned + tensor._vllm_warmup_ptr_range_bytes = value.ptr_range_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 766950f3321c..fa38e8a86a28 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_triton_helper import ( + DispatchSpec, + TritonWarmupTensor, + triton_kernel_dispatcher_with_warmup, + triton_ptr_range, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -68,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, @@ -191,6 +199,81 @@ 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) + buf_ptr_range_bytes = triton_ptr_range(kv_cache) + # 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) + return dict( + 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) + ), + ape_ptr=TritonWarmupTensor(torch.float32, shape=(pool_size, head_dim)), + loc_ptr=TritonWarmupTensor(torch.int64), + write_mask_ptr=TritonWarmupTensor(torch.bool), + 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, + 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 +344,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 +361,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 +371,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: @@ -316,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, @@ -362,6 +436,56 @@ 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] + tail_ptr_range_bytes = triton_ptr_range(tail_kv_cache) + return dict( + key_ptr=TritonWarmupTensor(torch.bfloat16), + score_ptr=TritonWarmupTensor(torch.bfloat16), + tslot_ptr=TritonWarmupTensor(torch.int64), + tail_ptr=TritonWarmupTensor( + torch.bfloat16, ptr_range_bytes=tail_ptr_range_bytes + ), + n_tokens=1, + 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 +500,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 7730265685ff..9b9920852f9f 100644 --- a/vllm/models/glm5next/amd/sparse_indexer.py +++ b/vllm/models/glm5next/amd/sparse_indexer.py @@ -754,6 +754,22 @@ def __init__( raise NotImplementedError( "SparseAttnIndexerKpool on ROCm does not support DCP." ) + 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( + 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 f7103e75fd45..3d2febe95761 100644 --- a/vllm/models/glm5next/common/kda.py +++ b/vllm/models/glm5next/common/kda.py @@ -48,6 +48,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 @@ -358,6 +359,11 @@ def _a_log_weight_loader(param, loaded_weight): vllm_config.model_config.dtype, self.kda_lower_bound, ) + if self.kda_prefill_backend == "triton": + 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 4827df2d07dd..864f9da28fdb 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=[ @@ -97,6 +106,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 ): @@ -111,23 +150,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)