@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ed582b6a4c
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 03:45 UTC

[Aiter][ROCm] gdn_linear_attn kernel fusion (#40711)

Signed-off-by: Tres Popp <tres.popp@amd.com> Signed-off-by: Chuan Li <chuali@amd.com> Co-authored-by: hellozhuo <zhuo.su@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Tres committed May 8, 2026 at 08:11 UTC ed582b6a4cae6f58c952768716a45b92a8ac1f3f
4 files changed +448 -52
vllm/_aiter_ops.py
+19
@@ -1407,6 +1407,25 @@ class rocm_aiter_ops:
1407 # https://github.com/ROCm/aiter/blob/6a0e7b26ccf33164785531212cc2ec2cde0b9243/aiter/dist/device_communicators/custom_all_reduce.py#L272-L273
1408 return int(cls._ALL_REDUCE_MAX_SIZE / 2)
1409
1410 + @classmethod
1411 + @if_aiter_supported
1412 + def are_gdn_triton_kernels_available(cls) -> bool:
1413 + """Check if AITER Triton kernels for GDN attention are importable.
1414 +
1415 + These are optional Triton kernels (conv1d fast-path, gated delta net)
1416 + used by GatedDeltaNetAttention's decode fast-path. They may be absent
1417 + in older aiter builds.
1418 + """
1419 + if not cls._AITER_ENABLED:
1420 + return False
1421 + try:
1422 + import aiter.ops.triton.causal_conv1d_update_single_token # noqa: F401
1423 + import aiter.ops.triton.gated_delta_net # noqa: F401
1424 +
1425 + return True
1426 + except (ImportError, ModuleNotFoundError):
1427 + return False
1428 +
1429 @staticmethod
1430 @if_aiter_supported
1431 def register_ops_once() -> None:
vllm/model_executor/layers/fla/ops/chunk.py
+11
@@ -32,6 +32,7 @@ def chunk_gated_delta_rule_fwd(
32 cu_seqlens: torch.Tensor | None = None,
33 chunk_indices: torch.Tensor | None = None,
34 chunk_offsets: torch.Tensor | None = None,
35 + core_attn_out: torch.Tensor | None = None,
36 ):
37 g = chunk_local_cumsum(
38 g, chunk_size=FLA_CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices
@@ -77,6 +78,7 @@ def chunk_gated_delta_rule_fwd(
78 scale=scale,
79 cu_seqlens=cu_seqlens,
80 chunk_indices=chunk_indices,
81 + core_attn_out=core_attn_out,
82 )
83 if SUPPRESS_LEVEL < 3:
84 return g, o, A, final_state, None, None, None
@@ -102,6 +104,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
104 chunk_indices: torch.Tensor | None = None,
105 chunk_offsets: torch.Tensor | None = None,
106 use_qk_l2norm_in_kernel: bool = False,
107 + core_attn_out: torch.Tensor | None = None,
108 ):
109 if use_qk_l2norm_in_kernel:
110 q = l2norm_fwd(q)
@@ -119,9 +122,15 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
122 cu_seqlens=cu_seqlens,
123 chunk_indices=chunk_indices,
124 chunk_offsets=chunk_offsets,
125 + core_attn_out=core_attn_out,
126 )
127 ctx.scale = scale
128 ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
129 + if core_attn_out is not None:
130 + assert not torch.is_grad_enabled(), (
131 + "core_attn_out buffer reuse is only supported for inference"
132 + )
133 + assert q.dtype == o.dtype, "Incompatible dtype for inplace computation"
134 return o.to(q.dtype), final_state
135
136
@@ -139,6 +148,7 @@ def chunk_gated_delta_rule(
148 chunk_indices: torch.Tensor | None = None,
149 chunk_offsets: torch.Tensor | None = None,
150 use_qk_l2norm_in_kernel: bool = False,
151 + core_attn_out: torch.Tensor | None = None,
152 ):
153 r"""
154 Args:
@@ -230,5 +240,6 @@ def chunk_gated_delta_rule(
240 chunk_indices,
241 chunk_offsets,
242 use_qk_l2norm_in_kernel,
243 + core_attn_out,
244 )
245 return o, final_state
vllm/model_executor/layers/fla/ops/chunk_o.py
+8 -1
@@ -148,6 +148,7 @@ def chunk_fwd_o(
148 cu_seqlens: torch.Tensor | None = None,
149 chunk_indices: torch.Tensor | None = None,
150 chunk_size: int = FLA_CHUNK_SIZE,
151 + core_attn_out: torch.Tensor | None = None,
152 ) -> torch.Tensor:
153 B, T, Hg, K, V = *q.shape, v.shape[-1]
154 H = v.shape[-2]
@@ -158,7 +159,13 @@ def chunk_fwd_o(
159 if scale is None:
160 scale = k.shape[-1] ** -0.5
161
161 - o = torch.empty_like(v)
162 + if core_attn_out is not None:
163 + assert core_attn_out.numel() >= v.numel(), (
164 + f"core_attn_out too small: {core_attn_out.numel()} < {v.numel()}"
165 + )
166 + o = core_attn_out[: v.numel()].view(*v.shape)
167 + else:
168 + o = torch.empty_like(v)
169
170 def grid(meta):
171 return (triton.cdiv(V, meta["BV"]), NT, B * H)
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+410 -51
@@ -8,6 +8,7 @@ from torch import nn
8 from transformers.activations import ACT2FN
9
10 from vllm import envs
11 +from vllm._aiter_ops import rocm_aiter_ops
12 from vllm.config import (
13 VllmConfig,
14 get_current_vllm_config,
@@ -64,6 +65,20 @@ from vllm.utils.torch_utils import (
65 )
66 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
67
68 +# Optional ROCm AITER Triton kernels for the GDN decode fast-path.
69 +# Availability is checked centrally via rocm_aiter_ops; the actual function
70 +# references are imported here so that they can be called without per-call
71 +# import overhead.
72 +GDN_AITER_TRITON_AVAILABLE = rocm_aiter_ops.are_gdn_triton_kernels_available()
73 +
74 +if GDN_AITER_TRITON_AVAILABLE:
75 + from aiter.ops.triton.causal_conv1d_update_single_token import (
76 + fused_reshape_causal_conv1d_update_single_token as gdn_aiter_fused_reshape_causal_conv1d_update_single_token, # noqa: E501
77 + )
78 + from aiter.ops.triton.gated_delta_net.fused_rearrange_sigmoid_gdr import (
79 + fused_rearrange_sigmoid_gated_delta_rule as gdn_aiter_fused_rearrange_sigmoid_gated_delta_rule, # noqa: E501
80 + )
81 +
82 logger = init_logger(__name__)
83
84
@@ -169,8 +184,9 @@ class ChunkGatedDeltaRule(CustomOp):
184 chunk_indices: torch.Tensor | None = None,
185 chunk_offsets: torch.Tensor | None = None,
186 use_qk_l2norm_in_kernel: bool = True,
187 + core_attn_out: torch.Tensor | None = None,
188 ):
173 - return fi_chunk_gated_delta_rule(
189 + o, final_state = fi_chunk_gated_delta_rule(
190 q=q,
191 k=k,
192 v=v,
@@ -181,6 +197,11 @@ class ChunkGatedDeltaRule(CustomOp):
197 cu_seqlens=cu_seqlens,
198 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
199 )
200 + if core_attn_out is not None:
201 + o_flat = o.squeeze(0).reshape(-1)
202 + co_flat = core_attn_out.reshape(-1)
203 + co_flat[: o_flat.numel()].copy_(o_flat)
204 + return o, final_state
205
206 def forward_native(
207 self,
@@ -195,6 +216,7 @@ class ChunkGatedDeltaRule(CustomOp):
216 chunk_indices: torch.Tensor | None = None,
217 chunk_offsets: torch.Tensor | None = None,
218 use_qk_l2norm_in_kernel: bool = True,
219 + core_attn_out: torch.Tensor | None = None,
220 ):
221 return fla_chunk_gated_delta_rule(
222 q=q,
@@ -208,6 +230,7 @@ class ChunkGatedDeltaRule(CustomOp):
230 chunk_indices=chunk_indices,
231 chunk_offsets=chunk_offsets,
232 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
233 + core_attn_out=core_attn_out,
234 )
235
236
@@ -271,7 +294,6 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
294 else 0
295 )
296 self.gqa_interleaved_layout = gqa_interleaved_layout
274 - self._forward_method = self.forward_cuda
297 if current_platform.is_xpu():
298 self._forward_method = self.forward_xpu
299 elif current_platform.is_cpu():
@@ -281,6 +303,10 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
303
304 register_cpu_gdn_attention_ops()
305 self._forward_method = self.forward_cpu
306 + elif current_platform.is_rocm():
307 + self._forward_method = self.forward_hip
308 + else:
309 + self._forward_method = self.forward_cuda
310
311 # QKV
312 self.conv_dim = self.key_dim * 2 + self.value_dim
@@ -297,6 +323,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
323 # we need to create qkvz_proj adaptively here.
324 # When create_in_proj_qkvz is False (e.g. LoRA enabled in Qwen3.5),
325 # in_proj_qkv and in_proj_z are created separately instead.
326 + self.has_lora_projections = not create_in_proj_qkvz
327 if create_in_proj_qkvz:
328 self.in_proj_qkvz = self.create_qkvz_proj(
329 hidden_size=self.hidden_size,
@@ -497,24 +524,155 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
524
525 return query, key, value, z, b, a
526
500 - def rearrange_mixed_qkv(self, mixed_qkv):
501 - if mixed_qkv is None:
502 - return None, None, None
503 - query, key, value = torch.split(
504 - mixed_qkv,
527 + @torch.compile(fullgraph=True)
528 + def prepare_gdn_attention_core_inputs(
529 + self,
530 + mixed_qkvz: torch.Tensor,
531 + mixed_ba: torch.Tensor,
532 + num_tokens: int,
533 + ):
534 + """
535 + Derives mixed_qkv, z, b, a from projected qkvz/ba for the GDN custom op.
536 +
537 + For gqa_interleaved_layout (Qwen3-Next): unpack the interleaved
538 + [ng, (hk + hk + np/ng*hv + np/ng*hv)] layout into contiguous qkv.
539 + For non-interleaved layout (Qwen3.5): simple split along last dim.
540 + """
541 + if not self.gqa_interleaved_layout:
542 + # Qwen3.5: weights are in [q, k, v, z] order
543 + assert num_tokens == mixed_qkvz.shape[0]
544 + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
545 + z_size = self.value_dim // self.tp_size
546 + mixed_qkv, z_flat = mixed_qkvz.split([qkv_size, z_size], dim=-1)
547 + n = mixed_qkvz.shape[0]
548 + z_out = z_flat.reshape(n, -1, self.head_v_dim)
549 + b, a = mixed_ba.chunk(2, dim=-1)
550 + return mixed_qkv, z_out, b, a
551 +
552 + # Qwen3-Next: interleaved GQA layout
553 + base_shape_qkvz = mixed_qkvz.size()[:-1]
554 + base_shape_ba = mixed_ba.size()[:-1]
555 + ng = self.num_k_heads // self.tp_size
556 +
557 + new_tensor_shape_qkvz = base_shape_qkvz + (
558 + ng,
559 + (
560 + self.head_k_dim
561 + + self.head_k_dim
562 + + (self.head_v_dim + self.head_v_dim)
563 + * self.num_v_heads
564 + // self.num_k_heads
565 + ),
566 + )
567 + new_tensor_shape_ba = base_shape_ba + (
568 + ng,
569 + 2 * self.num_v_heads // self.num_k_heads,
570 + )
571 +
572 + mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz)
573 + mixed_ba = mixed_ba.view(*new_tensor_shape_ba)
574 +
575 + split_arg_list_qkvz = [
576 + self.head_k_dim,
577 + self.head_k_dim,
578 + (self.num_v_heads // self.num_k_heads * self.head_v_dim),
579 + (self.num_v_heads // self.num_k_heads * self.head_v_dim),
580 + ]
581 + split_arg_list_ba = [
582 + self.num_v_heads // self.num_k_heads,
583 + self.num_v_heads // self.num_k_heads,
584 + ]
585 +
586 + (query, key, value, z) = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=-1)
587 + (b, a) = torch.split(mixed_ba, split_arg_list_ba, dim=-1)
588 +
589 + mixed_qkv_logical = torch.cat(
590 [
506 - self.key_dim // self.tp_size,
507 - self.key_dim // self.tp_size,
508 - self.value_dim // self.tp_size,
591 + query.reshape(num_tokens, -1),
592 + key.reshape(num_tokens, -1),
593 + value.reshape(num_tokens, -1),
594 ],
595 dim=-1,
596 )
512 - query, key = map(
513 - lambda x: rearrange(x, "l (h d) -> 1 l h d", d=self.head_k_dim),
514 - (query, key),
597 +
598 + # The split above produces non-contiguous views into the interleaved
599 + # buffer. Concatenating everything into a single flat tensor forces a
600 + # contiguous copy, then slicing back out gives contiguous q/k/v/z/b/a
601 + # tensors that downstream kernels require. Doing this in one cat+slice
602 + # keeps torch.compile in a single Triton graph instead of emitting
603 + # separate copy kernels per tensor. The original code used
604 + # rearrange(...).contiguous() on each tensor individually.
605 + fused = torch.cat(
606 + [
607 + mixed_qkv_logical.reshape(-1),
608 + z.reshape(-1),
609 + b.reshape(-1),
610 + a.reshape(-1),
611 + ],
612 + dim=0,
613 + )
614 +
615 + curr = 0
616 + qkv_numel = mixed_qkv_logical.numel()
617 + z_numel = z.numel()
618 + b_numel = b.numel()
619 + a_numel = a.numel()
620 +
621 + mixed_qkv_out = fused[curr : curr + qkv_numel].view(num_tokens, -1)
622 + curr += qkv_numel
623 +
624 + z_out = fused[curr : curr + z_numel].view(
625 + num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim
626 + )
627 + curr += z_numel
628 +
629 + b_out = fused[curr : curr + b_numel].view(
630 + num_tokens, self.num_v_heads // self.tp_size
631 + )
632 + curr += b_numel
633 +
634 + a_out = fused[curr : curr + a_numel].view(
635 + num_tokens, self.num_v_heads // self.tp_size
636 + )
637 +
638 + return mixed_qkv_out, z_out, b_out, a_out
639 +
640 + @torch.compile(fullgraph=True)
641 + def rearrange_mixed_qkv(self, mixed_qkv):
642 + """Split packed qkv into contiguous (1, seq, heads, dim) tensors.
643 +
644 + The original code used ``rearrange(x, "l (h d) -> 1 l h d", d=...)``
645 + followed by ``.contiguous()`` on each tensor. This version flattens
646 + all three splits into a single buffer via ``torch.cat`` so that
647 + torch.compile emits one Triton copy kernel instead of three separate
648 + contiguous() calls.
649 + """
650 + if mixed_qkv is None:
651 + return None, None, None
652 +
653 + seq_len = mixed_qkv.shape[0]
654 + q_dim = self.key_dim // self.tp_size
655 + k_dim = self.key_dim // self.tp_size
656 + v_dim = self.value_dim // self.tp_size
657 +
658 + query, key, value = torch.split(mixed_qkv, [q_dim, k_dim, v_dim], dim=-1)
659 +
660 + fused = torch.cat(
661 + [query.reshape(-1), key.reshape(-1), value.reshape(-1)], dim=0
662 )
516 - value = rearrange(value, "l (h d) -> 1 l h d", d=self.head_v_dim)
517 - return query.contiguous(), key.contiguous(), value.contiguous()
663 +
664 + q_size = seq_len * q_dim
665 + k_size = seq_len * k_dim
666 +
667 + q_contig = fused[0:q_size]
668 + k_contig = fused[q_size : q_size + k_size]
669 + v_contig = fused[q_size + k_size :]
670 +
671 + query = q_contig.view(1, seq_len, -1, self.head_k_dim)
672 + key = k_contig.view(1, seq_len, -1, self.head_k_dim)
673 + value = v_contig.view(1, seq_len, -1, self.head_v_dim)
674 +
675 + return query, key, value
676
677 def forward(
678 self,
@@ -523,6 +681,63 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
681 ):
682 self._forward_method(hidden_states, output)
683
684 + def _output_projection(
685 + self,
686 + core_attn_out: torch.Tensor,
687 + z: torch.Tensor,
688 + output: torch.Tensor,
689 + num_tokens: int,
690 + ):
691 + """Part 3: RMSNormGated + output linear projection.
692 +
693 + The RMSNormGated + quant sequence is eligible for fusion
694 + by the compilation pass when fuse_norm_quant is enabled.
695 + """
696 + z_shape_og = z.shape
697 + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
698 + z = z.reshape(-1, z.shape[-1])
699 + core_attn_out = self.norm(core_attn_out, z)
700 + core_attn_out = core_attn_out.reshape(z_shape_og)
701 + core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
702 + output[:num_tokens], _ = self.out_proj(core_attn_out)
703 +
704 + def forward_hip(
705 + self,
706 + hidden_states: torch.Tensor,
707 + output: torch.Tensor,
708 + ):
709 + """ROCm forward using AITER Triton fused projection+attention when
710 + available, otherwise falling back to the generic CUDA path."""
711 + if not self.has_lora_projections and GDN_AITER_TRITON_AVAILABLE:
712 + num_tokens = hidden_states.size(0)
713 + projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)
714 + projected_states_ba, _ = self.in_proj_ba(hidden_states)
715 + projected_states_qkvz = projected_states_qkvz.view(num_tokens, -1)
716 + projected_states_ba = projected_states_ba.view(num_tokens, -1)
717 + core_attn_out = torch.empty(
718 + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
719 + dtype=hidden_states.dtype,
720 + device=hidden_states.device,
721 + )
722 + z = torch.empty(
723 + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
724 + dtype=projected_states_qkvz.dtype,
725 + device=projected_states_qkvz.device,
726 + )
727 +
728 + torch.ops.vllm.gdn_attention_core(
729 + projected_states_qkvz,
730 + projected_states_ba,
731 + z,
732 + core_attn_out,
733 + fast_kernel=True,
734 + layer_name=_encode_layer_name(self.prefix),
735 + )
736 +
737 + self._output_projection(core_attn_out, z, output, num_tokens)
738 + else:
739 + self.forward_cuda(hidden_states, output)
740 +
741 def forward_cuda(
742 self,
743 hidden_states: torch.Tensor,
@@ -538,7 +753,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
753 # ============================================================
754 # Part 1: Input Projection
755 # ============================================================
541 - if hasattr(self, "in_proj_qkv"):
756 + if self.has_lora_projections:
757 # LoRA path (Qwen3.5 only): separate in_proj_qkv and in_proj_z
758 mixed_qkv, _ = self.in_proj_qkv(hidden_states)
759 ba, _ = self.in_proj_ba(hidden_states)
@@ -586,20 +801,14 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
801 b,
802 a,
803 core_attn_out,
589 - _encode_layer_name(self.prefix),
804 + fast_kernel=False,
805 + layer_name=_encode_layer_name(self.prefix),
806 )
807
808 # ============================================================
809 # Part 3: Output Projection
810 # ============================================================
595 - z_shape_og = z.shape
596 - # Reshape input data into 2D tensor
597 - core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
598 - z = z.reshape(-1, z.shape[-1])
599 - core_attn_out = self.norm(core_attn_out, z)
600 - core_attn_out = core_attn_out.reshape(z_shape_og)
601 - core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
602 - output[:num_tokens], _ = self.out_proj(core_attn_out)
811 + self._output_projection(core_attn_out, z, output, num_tokens)
812
813 def forward_xpu(
814 self,
@@ -614,7 +823,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
823 """
824 num_tokens = hidden_states.size(0)
825
617 - assert not hasattr(self, "in_proj_qkv"), "lora isn't supported on XPU."
826 + assert not self.has_lora_projections, "lora isn't supported on XPU."
827
828 # ============================================================
829 # Part 1: Input Projection
@@ -702,7 +911,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
911 core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
912 output[:num_tokens], _ = self.out_proj(core_attn_out)
913
705 - def _warmup_prefill_kernels(self, mixed_qkv: torch.Tensor) -> None:
914 + def _warmup_prefill_kernels(self, qkv_or_qkvz: torch.Tensor, v_dim: int) -> None:
915 """Warm up GDN prefill kernels during V1 profiling.
916
917 During V1 profile runs, ``_forward_core`` returns early because
@@ -723,7 +932,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
932 ``BT = chunk_size`` (64). A single warmup pass with T = 64
933 is sufficient to populate the autotuner cache.
934
726 - The decode path uses ``fused_sigmoid_gating_delta_rule_update``
935 + The decode path uses ``gdn_aiter_fused_rearrange_sigmoid_gated_delta_rule``
936 which has fixed kernel parameters (no autotuning), so only the
937 prefill (chunked) path needs warming up.
938 """
@@ -731,8 +940,8 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
940 return
941 self._prefill_kernels_warmed_up = True
942
734 - device = mixed_qkv.device
735 - dtype = mixed_qkv.dtype
943 + device = qkv_or_qkvz.device
944 + dtype = qkv_or_qkvz.dtype
945 num_k_heads = self.num_k_heads // self.tp_size
946 num_v_heads = self.num_v_heads // self.tp_size
947 _, state_dtype = self.get_state_dtype()
@@ -743,7 +952,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
952 # then run chunk_gated_delta_rule with in-kernel L2 norm disabled.
953 T = FLA_CHUNK_SIZE
954 dummy_mixed_qkv = torch.randn(
746 - T, mixed_qkv.shape[-1], device=device, dtype=dtype
955 + T, qkv_or_qkvz.shape[-1] - v_dim, device=device, dtype=dtype
956 )
957 dummy_a = torch.randn(T, num_v_heads, device=device, dtype=dtype)
958 dummy_b = torch.randn(T, num_v_heads, device=device, dtype=dtype)
@@ -806,6 +1015,66 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
1015
1016 torch.accelerator.empty_cache()
1017
1018 + def _forward_core_rocm(
1019 + self,
1020 + qkvz: torch.Tensor,
1021 + ba: torch.Tensor,
1022 + z_out: torch.Tensor,
1023 + core_attn_out: torch.Tensor,
1024 + ):
1025 + """ROCm AITER fast path: conv1d + recurrent attention from packed
1026 + qkvz/ba layout.
1027 +
1028 + For decode-only (no spec, no prefill), dispatches directly to
1029 + ``_forward_core_decode_fast``. Otherwise unpacks the packed
1030 + layout and falls through to ``_forward_core``.
1031 +
1032 + Args:
1033 + qkvz: packed [q, k, v, z] projection (num_tokens, qkvz_dim)
1034 + ba: packed [b, a] gating vectors (num_tokens, 2*num_heads)
1035 + z_out: **output** buffer for z (num_tokens, num_heads,
1036 + head_dim); mutated in-place.
1037 + core_attn_out: Pre-allocated output buffer for attention results.
1038 + """
1039 + forward_context = get_forward_context()
1040 + attn_metadata_raw = forward_context.attn_metadata
1041 +
1042 + if attn_metadata_raw is None:
1043 + v_dim = core_attn_out.shape[-1] * core_attn_out.shape[-2]
1044 + self._warmup_prefill_kernels(qkvz, v_dim)
1045 + return
1046 +
1047 + assert isinstance(attn_metadata_raw, dict)
1048 + attn_metadata = attn_metadata_raw[self.prefix] # type: ignore[index]
1049 + assert isinstance(attn_metadata, GDNAttentionMetadata)
1050 +
1051 + if (
1052 + attn_metadata.spec_sequence_masks is None
1053 + and attn_metadata.num_prefills == 0
1054 + and attn_metadata.num_decodes > 0
1055 + ):
1056 + return self._forward_core_decode_fast(
1057 + qkvz=qkvz,
1058 + ba=ba,
1059 + z_out=z_out,
1060 + core_attn_out=core_attn_out,
1061 + attn_metadata=attn_metadata,
1062 + )
1063 +
1064 + core_attn_out.zero_()
1065 + z_out.zero_()
1066 + num_tokens_all = qkvz.shape[0]
1067 + mixed_qkv, z, b, a = self.prepare_gdn_attention_core_inputs(
1068 + qkvz, ba, num_tokens_all
1069 + )
1070 + z_out[:] = z
1071 + self._forward_core(
1072 + mixed_qkv=mixed_qkv,
1073 + b=b,
1074 + a=a,
1075 + core_attn_out=core_attn_out,
1076 + )
1077 +
1078 def _forward_core(
1079 self,
1080 mixed_qkv: torch.Tensor,
@@ -813,13 +1082,19 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
1082 a: torch.Tensor,
1083 core_attn_out: torch.Tensor,
1084 ):
1085 + """Core conv1d + recurrent attention (standard path).
1086 +
1087 + Args:
1088 + mixed_qkv: packed [q, k, v] projection (num_tokens, qkv_dim)
1089 + b: beta gating vector (num_tokens, num_heads)
1090 + a: alpha gating vector (num_tokens, num_heads)
1091 + core_attn_out: Pre-allocated output buffer for attention results.
1092 + """
1093 forward_context = get_forward_context()
1094 attn_metadata_raw = forward_context.attn_metadata
1095
1096 if attn_metadata_raw is None:
820 - # V1 profile run — warm up prefill kernels so that
821 - # autotuning completes before KV cache allocation.
822 - self._warmup_prefill_kernels(mixed_qkv)
1097 + self._warmup_prefill_kernels(mixed_qkv, 0)
1098 return
1099
1100 assert isinstance(attn_metadata_raw, dict)
@@ -1065,6 +1340,72 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
1340 else:
1341 core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
1342
1343 + def _forward_core_decode_fast(
1344 + self,
1345 + qkvz: torch.Tensor,
1346 + ba: torch.Tensor,
1347 + z_out: torch.Tensor,
1348 + core_attn_out: torch.Tensor,
1349 + attn_metadata: GDNAttentionMetadata,
1350 + ):
1351 + non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
1352 + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
1353 + self_kv_cache = self.kv_cache
1354 + # conv_state must be (..., dim, width-1) for the conv kernels.
1355 + # DS layout stores it that way directly; SD layout needs a transpose.
1356 + conv_state = (
1357 + self_kv_cache[0]
1358 + if is_conv_state_dim_first()
1359 + else self_kv_cache[0].transpose(-1, -2)
1360 + )
1361 + ssm_state = self_kv_cache[1]
1362 +
1363 + # 1. Convolution sequence transformation
1364 + conv_weights = self.conv1d.weight.view(
1365 + self.conv1d.weight.size(0), self.conv1d.weight.size(2)
1366 + )
1367 +
1368 + mixed_qkv_non_spec, b, a = (
1369 + gdn_aiter_fused_reshape_causal_conv1d_update_single_token(
1370 + qkvz,
1371 + attn_metadata.num_actual_tokens,
1372 + self.num_k_heads // self.tp_size,
1373 + self.num_v_heads // self.tp_size,
1374 + self.head_k_dim,
1375 + self.head_v_dim,
1376 + ba,
1377 + z_out,
1378 + core_attn_out,
1379 + conv_state,
1380 + conv_weights,
1381 + self.conv1d.bias,
1382 + self.activation,
1383 + conv_state_indices=non_spec_state_indices_tensor[ # type: ignore[index]
1384 + : attn_metadata.num_actual_tokens
1385 + ],
1386 + validate_data=True,
1387 + )
1388 + )
1389 +
1390 + # 2. Recurrent attention
1391 + gdn_aiter_fused_rearrange_sigmoid_gated_delta_rule(
1392 + A_log=self.A_log,
1393 + a=a,
1394 + b=b,
1395 + dt_bias=self.dt_bias,
1396 + qkv=mixed_qkv_non_spec,
1397 + key_dim=self.key_dim // self.tp_size,
1398 + value_dim=self.value_dim // self.tp_size,
1399 + head_k_dim=self.head_k_dim,
1400 + head_v_dim=self.head_v_dim,
1401 + initial_state=ssm_state,
1402 + inplace_final_state=True,
1403 + cu_seqlens=non_spec_query_start_loc[: attn_metadata.num_decodes + 1], # type: ignore[index]
1404 + ssm_state_indices=non_spec_state_indices_tensor,
1405 + use_qk_l2norm_in_kernel=True,
1406 + core_attn_out=core_attn_out.reshape(-1),
1407 + )
1408 +
1409 def _forward_core_decode_non_spec(
1410 self,
1411 mixed_qkv: torch.Tensor,
@@ -1121,33 +1462,51 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
1462
1463
1464 def gdn_attention_core(
1124 - mixed_qkv: torch.Tensor,
1125 - b: torch.Tensor,
1126 - a: torch.Tensor,
1465 + qkv_or_qkvz: torch.Tensor,
1466 + b_or_ba: torch.Tensor,
1467 + a_or_z_out: torch.Tensor,
1468 core_attn_out: torch.Tensor,
1469 + fast_kernel: bool,
1470 layer_name: LayerNameType,
1471 ) -> None:
1130 - """
1131 - Custom op for the core attention computation.
1132 - Only handles the convolution + recurrent attention part.
1133 - Input/output projections are handled outside this op.
1472 + """Custom op dispatching to _forward_core or _forward_core_rocm.
1473 +
1474 + Handles conv1d + recurrent attention only; input/output projections
1475 + are performed by the caller.
1476 +
1477 + When ``fast_kernel=False`` (standard path):
1478 + qkv_or_qkvz is [q, k, v], b_or_ba is b, a_or_z_out is a (read-only).
1479 + When ``fast_kernel=True`` (AITER Triton fast path, ROCm only):
1480 + qkv_or_qkvz is [q, k, v, z], b_or_ba is [b, a], a_or_z_out is the
1481 + z output buffer (mutated in-place).
1482 +
1483 + ``core_attn_out`` is always mutated in-place.
1484 """
1485 layer_name = _resolve_layer_name(layer_name)
1486 forward_context: ForwardContext = get_forward_context()
1487 self = forward_context.no_compile_layers[layer_name]
1138 - self._forward_core(
1139 - mixed_qkv=mixed_qkv,
1140 - b=b,
1141 - a=a,
1142 - core_attn_out=core_attn_out,
1143 - )
1488 + if fast_kernel:
1489 + self._forward_core_rocm(
1490 + qkvz=qkv_or_qkvz,
1491 + ba=b_or_ba,
1492 + z_out=a_or_z_out,
1493 + core_attn_out=core_attn_out,
1494 + )
1495 + else:
1496 + self._forward_core(
1497 + mixed_qkv=qkv_or_qkvz,
1498 + b=b_or_ba,
1499 + a=a_or_z_out,
1500 + core_attn_out=core_attn_out,
1501 + )
1502
1503
1504 def gdn_attention_core_fake(
1147 - mixed_qkv: torch.Tensor,
1148 - b: torch.Tensor,
1149 - a: torch.Tensor,
1505 + qkv_or_qkvz: torch.Tensor,
1506 + b_or_ba: torch.Tensor,
1507 + a_or_z_out: torch.Tensor,
1508 core_attn_out: torch.Tensor,
1509 + fast_kernel: bool,
1510 layer_name: LayerNameType,
1511 ) -> None:
1512 """Fake implementation for torch.compile."""
@@ -1157,7 +1516,7 @@ def gdn_attention_core_fake(
1516 direct_register_custom_op(
1517 op_name="gdn_attention_core",
1518 op_func=gdn_attention_core,
1160 - mutates_args=["core_attn_out"],
1519 + mutates_args=["a_or_z_out", "core_attn_out"],
1520 fake_impl=gdn_attention_core_fake,
1521 )
1522