Skip to content

vllm.model_executor.layers.mamba.ops.hip_gdn_chunk

Gated delta net prefill on RDNA3.5, as a single HIP kernel.

Replaces the cumsum, WY transform, state recurrence and output kernels with one launch of torch.ops._rocm_C.rdna35_gdn_chunked.

Functions:

chunk_gdn_hip_fwd(q, k, v, g, beta, scale, initial_state, cu_seqlens, core_attn_out=None)

Run the gated delta rule and return (output, final_state).

g is the raw per-token log decay; the kernel takes the cumsum itself, so this must be called before any chunk-local cumsum.

Source code in vllm/model_executor/layers/mamba/ops/hip_gdn_chunk.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def chunk_gdn_hip_fwd(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    scale: float,
    initial_state: torch.Tensor | None,
    cu_seqlens: torch.Tensor,
    core_attn_out: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Run the gated delta rule and return ``(output, final_state)``.

    ``g`` is the raw per-token log decay; the kernel takes the cumsum itself,
    so this must be called before any chunk-local cumsum.
    """
    head_key_dim = k.shape[-1]
    num_value_heads, head_value_dim = v.shape[-2], v.shape[-1]
    num_seqs = len(cu_seqlens) - 1

    if core_attn_out is not None:
        output = core_attn_out[: v.numel()].view(*v.shape)
    else:
        output = torch.empty_like(v)
    final_state = q.new_empty(
        num_seqs, num_value_heads, head_value_dim, head_key_dim, dtype=torch.float32
    )

    torch.ops._rocm_C.rdna35_gdn_chunked(
        q.squeeze(0),
        k.squeeze(0),
        v.squeeze(0),
        g.squeeze(0).float(),
        beta.squeeze(0).float(),
        None if initial_state is None else initial_state.to(torch.float32).contiguous(),
        cu_seqlens.to(torch.int32),
        output.squeeze(0),
        final_state,
        float(scale),
    )
    return output, final_state

is_hip_gdn_available() cached

Whether the op is in this build and the device can run it.

The device test is the RDNA3.5 family rather than a broader question such as "is this Navi": gfx1100, gfx1103 and gfx12xx would answer yes, and the kernel's wave32 block layout does not hold on them.

Source code in vllm/model_executor/layers/mamba/ops/hip_gdn_chunk.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
@functools.cache
def is_hip_gdn_available() -> bool:
    """Whether the op is in this build and the device can run it.

    The device test is the RDNA3.5 family rather than a broader question such
    as "is this Navi": gfx1100, gfx1103 and gfx12xx would answer yes, and the
    kernel's wave32 block layout does not hold on them.
    """
    if not hasattr(torch.ops, "_rocm_C") or not hasattr(
        torch.ops._rocm_C, "rdna35_gdn_chunked"
    ):
        return False
    # Imported here so a non-ROCm platform never pays for it.
    from vllm.platforms.rocm import on_gfx115x

    return on_gfx115x()

is_hip_gdn_supported(query, value, cu_seqlens)

Whether :func:chunk_gdn_hip_fwd can serve this call.

Mirrors every shape and dtype the kernel asserts on, so an unsupported call falls back to the Triton kernels instead of raising.

Source code in vllm/model_executor/layers/mamba/ops/hip_gdn_chunk.py 
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def is_hip_gdn_supported(
    query: torch.Tensor,
    value: torch.Tensor,
    cu_seqlens: torch.Tensor | None,
) -> bool:
    """Whether :func:`chunk_gdn_hip_fwd` can serve this call.

    Mirrors every shape and dtype the kernel asserts on, so an unsupported
    call falls back to the Triton kernels instead of raising.
    """
    if not is_hip_gdn_available():
        return False
    if cu_seqlens is None or query.shape[0] != 1:
        return False
    if query.dtype not in (torch.bfloat16, torch.float16):
        return False
    if value.dtype != query.dtype:
        return False
    if query.shape[-1] != 128 or value.shape[-1] != 128:
        return False
    num_key_heads, num_value_heads = query.shape[-2], value.shape[-2]
    return num_value_heads % num_key_heads == 0