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 5, 17:15 UTC
[Model] Sync upstream BT=chunk_size fix for GDN chunk_fwd_kernel_o, simplify warmup to single pass (#38343)
Signed-off-by: AuYang <459461160@qq.com> Co-authored-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Xu Jinyang committed
Apr 1, 2026 at 03:03 UTC
b779eb3363193c78bad71e4a32db6563270aa096
3 files changed
+59
-66
vllm/model_executor/layers/fla/ops/chunk_o.py
+3
-3
@@ -16,7 +16,7 @@ from vllm.triton_utils import tl, triton
16
17
from .index import prepare_chunk_indices
18
from .op import exp
19
-from .utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper
19
+from .utils import FLA_CHUNK_SIZE, check_shared_mem, is_nvidia_hopper
20
21
BKV_LIST = [64, 128] if check_shared_mem() else [32, 64]
22
NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8]
@@ -146,11 +146,11 @@ def chunk_fwd_o(
146
g: torch.Tensor | None = None, # cumsum of log decay
147
scale: float | None = None,
148
cu_seqlens: torch.Tensor | None = None,
149
- chunk_size: int = 64,
149
+ chunk_size: int = FLA_CHUNK_SIZE,
150
) -> torch.Tensor:
151
B, T, Hg, K, V = *q.shape, v.shape[-1]
152
H = v.shape[-2]
153
- BT = 64 if FLA_GDN_FIX_BT else min(chunk_size, max(16, triton.next_power_of_2(T)))
153
+ BT = chunk_size
154
chunk_indices = (
155
prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
156
)
vllm/model_executor/layers/fla/ops/utils.py
+3
-1
@@ -24,10 +24,12 @@ logger = logging.getLogger(__name__)
24
25
COMPILER_MODE = os.getenv("FLA_COMPILER_MODE") == "1"
26
FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1"
27
-FLA_GDN_FIX_BT = os.getenv("FLA_GDN_FIX_BT", "0") == "1"
27
28
SUPPRESS_LEVEL = int(os.getenv("GDN_RECOMPUTE_SUPPRESS_LEVEL", "0"))
29
30
+# Default chunk size used across FLA triton kernels (kda, chunk, chunk_o, etc.)
31
+FLA_CHUNK_SIZE = 64
32
+
33
34
def tensor_cache(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:
35
"""
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+53
-62
@@ -28,6 +28,7 @@ from vllm.model_executor.layers.fla.ops import (
28
fused_sigmoid_gating_delta_rule_update,
29
)
30
from vllm.model_executor.layers.fla.ops.chunk import l2norm_fwd
31
+from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE
32
from vllm.model_executor.layers.layernorm import RMSNormGated
33
from vllm.model_executor.layers.linear import (
34
ColumnParallelLinear,
@@ -581,11 +582,9 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
582
results are cached globally, so only the first layer incurs
583
actual benchmarking cost.
584
584
- Most kernels use a fixed ``BT = chunk_size`` (64), but
585
- ``chunk_fwd_kernel_o`` recomputes ``BT`` from the sequence
586
- length: ``min(64, max(16, next_power_of_2(T)))``. Since ``BT``
587
- is part of its autotune key, we run warmup passes with T = 16,
588
- 32, and 64 to cover all possible ``BT`` values.
585
+ All kernels including ``chunk_fwd_kernel_o`` now use a fixed
586
+ ``BT = chunk_size`` (64). A single warmup pass with T = 64
587
+ is sufficient to populate the autotuner cache.
588
589
The decode path uses ``fused_sigmoid_gating_delta_rule_update``
590
which has fixed kernel parameters (no autotuning), so only the
@@ -601,66 +600,58 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
600
num_v_heads = self.num_v_heads // self.tp_size
601
_, state_dtype = self.get_state_dtype()
602
604
- # Run warmup for each possible BT value of chunk_fwd_kernel_o:
605
- # T=16 → BT=16, T=32 → BT=32, T=64 → BT=64.
606
- # Other kernels always use BT=chunk_size(64), so their autotune
607
- # cache is populated on the first pass and reused thereafter.
608
- for T in (16, 32, 64):
609
- q = torch.randn(
610
- 1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype
611
- )
612
- k = torch.randn(
613
- 1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype
603
+ # All kernels use BT = chunk_size (FLA_CHUNK_SIZE4), so a single pass with
604
+ # T = chunk_size is sufficient to populate every autotuner cache.
605
+ T = FLA_CHUNK_SIZE
606
+ q = torch.randn(1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype)
607
+ k = torch.randn(1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype)
608
+ v = torch.randn(1, T, num_v_heads, self.head_v_dim, device=device, dtype=dtype)
609
+ # NOTE: g and beta must have the same dtypes as during
610
+ # inference, so we construct them with the same function
611
+ # (fused_gdn_gating). dummy_a and dummy_b are throwaway
612
+ # inputs required by that function.
613
+ dummy_a = torch.randn(T, num_v_heads, device=device, dtype=dtype)
614
+ dummy_b = torch.randn(T, num_v_heads, device=device, dtype=dtype)
615
+ g, beta = fused_gdn_gating(self.A_log, dummy_a, dummy_b, self.dt_bias)
616
+ state = torch.zeros(
617
+ 1,
618
+ num_v_heads,
619
+ self.head_v_dim,
620
+ self.head_k_dim,
621
+ device=device,
622
+ dtype=state_dtype,
623
+ )
624
+ cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.int32)
625
+
626
+ try:
627
+ self.chunk_gated_delta_rule(
628
+ q=q,
629
+ k=k,
630
+ v=v,
631
+ g=g,
632
+ beta=beta,
633
+ initial_state=state,
634
+ output_final_state=True,
635
+ cu_seqlens=cu_seqlens,
636
+ use_qk_l2norm_in_kernel=True,
637
)
615
- v = torch.randn(
616
- 1, T, num_v_heads, self.head_v_dim, device=device, dtype=dtype
638
+ except Exception:
639
+ logger.warning(
640
+ "GDN prefill kernel warmup (T=%d) failed for "
641
+ "layer %s. First inference may OOM due to "
642
+ "autotuner.",
643
+ T,
644
+ self.prefix,
645
+ exc_info=True,
646
)
618
- # NOTE: g and beta must have the same dtypes as during
619
- # inference, so we construct them with the same function
620
- # (fused_gdn_gating). dummy_a and dummy_b are throwaway
621
- # inputs required by that function.
622
- dummy_a = torch.randn(T, num_v_heads, device=device, dtype=dtype)
623
- dummy_b = torch.randn(T, num_v_heads, device=device, dtype=dtype)
624
- g, beta = fused_gdn_gating(self.A_log, dummy_a, dummy_b, self.dt_bias)
625
- state = torch.zeros(
626
- 1,
627
- num_v_heads,
628
- self.head_v_dim,
629
- self.head_k_dim,
630
- device=device,
631
- dtype=state_dtype,
647
+ else:
648
+ logger.debug(
649
+ "GDN prefill kernel warmup (T=%d) completed for layer %s",
650
+ T,
651
+ self.prefix,
652
)
633
- cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.int32)
634
-
635
- try:
636
- self.chunk_gated_delta_rule(
637
- q=q,
638
- k=k,
639
- v=v,
640
- g=g,
641
- beta=beta,
642
- initial_state=state,
643
- output_final_state=True,
644
- cu_seqlens=cu_seqlens,
645
- use_qk_l2norm_in_kernel=True,
646
- )
647
- except Exception:
648
- logger.warning(
649
- "GDN prefill kernel warmup (T=%d) failed for "
650
- "layer %s. First inference may OOM due to "
651
- "autotuner.",
652
- T,
653
- self.prefix,
654
- exc_info=True,
655
- )
656
- else:
657
- logger.debug(
658
- "GDN prefill kernel warmup (T=%d) completed for layer %s",
659
- T,
660
- self.prefix,
661
- )
662
- finally:
663
- del q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens
653
+ finally:
654
+ del q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens
655
656
torch.accelerator.empty_cache()
657