[JIT][ROCm] Fully warmup GLM-5.3-Flash

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:

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,

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

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],
     )