@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a16133a0f1

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 6, 05:45 UTC

[Perf] [Bugfix] Fix Triton autotuning in inference for Qwen3.5 (#37338)

Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>

Artem Perevedentsev committed Mar 23, 2026 at 09:37 UTC a16133a0f12e8e406206be1bc8041727d9f1ea7d

10 files changed +40 -39

vllm/model_executor/layers/fla/ops/chunk.py

+5 -5

@@ -30,7 +30,7 @@ def chunk_gated_delta_rule_fwd(

30 scale: float,

31 initial_state: torch.Tensor,

32 output_final_state: bool,

33 - cu_seqlens: torch.LongTensor | None = None,

33 + cu_seqlens: torch.Tensor | None = None,

34 ):

35 g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens)

36 # obtain WY representation. u is actually the new v.

@@ -84,7 +84,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):

84 scale: float,

85 initial_state: torch.Tensor,

86 output_final_state: bool,

87 - cu_seqlens: torch.LongTensor | None = None,

87 + cu_seqlens: torch.Tensor | None = None,

88 use_qk_l2norm_in_kernel: bool = False,

89 ):

90 if use_qk_l2norm_in_kernel:

@@ -117,7 +117,7 @@ def chunk_gated_delta_rule(

117 scale: float = None,

118 initial_state: torch.Tensor = None,

119 output_final_state: bool = False,

120 - cu_seqlens: torch.LongTensor | None = None,

120 + cu_seqlens: torch.Tensor | None = None,

121 use_qk_l2norm_in_kernel: bool = False,

122 ):

123 r"""

@@ -141,7 +141,7 @@ def chunk_gated_delta_rule(

141 Default: `None`.

142 output_final_state (Optional[bool]):

143 Whether to output the final state of shape `[N, H, V, K]`. Default: `False`.

144 - cu_seqlens (torch.LongTensor):

144 + cu_seqlens (torch.Tensor):

145 Cumulative sequence lengths of shape `[N+1]` used for variable-length training,

146 consistent with the FlashAttention API.

147 Returns:

@@ -171,7 +171,7 @@ def chunk_gated_delta_rule(

171 # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required

172 >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))

173 # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected

174 - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)

174 + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.int32)

175 >>> o_var, ht_var = chunk_gated_delta_rule(

176 q, k, v, g, beta,

177 initial_state=h0,

vllm/model_executor/layers/fla/ops/chunk_delta_h.py

+1 -1

@@ -288,7 +288,7 @@ def chunk_gated_delta_rule_fwd_h(

288 output_final_state: bool = False,

289 chunk_size: int = 64, # SY: remove this argument and force chunk size 64?

290 save_new_value: bool = True,

291 - cu_seqlens: torch.LongTensor | None = None,

291 + cu_seqlens: torch.Tensor | None = None,

292 ) -> tuple[torch.Tensor, torch.Tensor]:

293 # This kernel is slightly different from fla to support Q/K with different head numbers.

294 # In fla, Q/K always have the same head number, so Hg is always equal to H.

vllm/model_executor/layers/fla/ops/chunk_o.py

+1 -1

@@ -145,7 +145,7 @@ def chunk_fwd_o(

145 h: torch.Tensor,

146 g: torch.Tensor | None = None, # cumsum of log decay

147 scale: float | None = None,

148 - cu_seqlens: torch.LongTensor | None = None,

148 + cu_seqlens: torch.Tensor | None = None,

149 chunk_size: int = 64,

150 ) -> torch.Tensor:

151 B, T, Hg, K, V = *q.shape, v.shape[-1]

vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py

+2 -2

@@ -102,7 +102,7 @@ def chunk_scaled_dot_kkt_fwd(

102 k: torch.Tensor,

103 g: torch.Tensor | None = None,

104 beta: torch.Tensor | None = None,

105 - cu_seqlens: torch.LongTensor | None = None,

105 + cu_seqlens: torch.Tensor | None = None,

106 chunk_size: int = 64,

107 output_dtype: torch.dtype = torch.float32,

108 ) -> torch.Tensor:

@@ -116,7 +116,7 @@ def chunk_scaled_dot_kkt_fwd(

116 The beta tensor of shape `[B, T, H]`.

117 g (torch.Tensor):

118 The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`.

119 - cu_seqlens (torch.LongTensor):

119 + cu_seqlens (torch.Tensor):

120 The cumulative sequence lengths of the input tensor.

121 Default: None

122 chunk_size (int):

vllm/model_executor/layers/fla/ops/fused_recurrent.py

+5 -5

@@ -184,7 +184,7 @@ def fused_recurrent_gated_delta_rule_fwd(

184 scale: float,

185 initial_state: torch.Tensor,

186 inplace_final_state: bool = True,

187 - cu_seqlens: torch.LongTensor | None = None,

187 + cu_seqlens: torch.Tensor | None = None,

188 ssm_state_indices: torch.Tensor | None = None,

189 num_accepted_tokens: torch.Tensor | None = None,

190 use_qk_l2norm_in_kernel: bool = False,

@@ -489,7 +489,7 @@ class FusedRecurrentFunction(torch.autograd.Function):

489 scale: float,

490 initial_state: torch.Tensor,

491 inplace_final_state: bool = True,

492 - cu_seqlens: torch.LongTensor | None = None,

492 + cu_seqlens: torch.Tensor | None = None,

493 ssm_state_indices: torch.Tensor | None = None,

494 num_accepted_tokens: torch.Tensor | None = None,

495 use_qk_l2norm_in_kernel: bool = False,

@@ -521,7 +521,7 @@ def fused_recurrent_gated_delta_rule(

521 scale: float = None,

522 initial_state: torch.Tensor = None,

523 inplace_final_state: bool = True,

524 - cu_seqlens: torch.LongTensor | None = None,

524 + cu_seqlens: torch.Tensor | None = None,

525 ssm_state_indices: torch.Tensor | None = None,

526 num_accepted_tokens: torch.Tensor | None = None,

527 use_qk_l2norm_in_kernel: bool = False,

@@ -549,7 +549,7 @@ def fused_recurrent_gated_delta_rule(

549 inplace_final_state: bool:

550 Whether to store the final state in-place to save memory.

551 Default: `True`.

552 - cu_seqlens (torch.LongTensor):

552 + cu_seqlens (torch.Tensor):

553 Cumulative sequence lengths of shape `[N+1]` used for variable-length training,

554 consistent with the FlashAttention API.

555 ssm_state_indices (Optional[torch.Tensor]):

@@ -583,7 +583,7 @@ def fused_recurrent_gated_delta_rule(

583 # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required

584 >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))

585 # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected

586 - >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)

586 + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.int32)

587 >>> o_var, ht_var = fused_gated_recurrent_delta_rule(

588 q, k, v, g, beta,

589 initial_state=h0,

vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py

+1 -1

@@ -191,7 +191,7 @@ def fused_sigmoid_gating_delta_rule_update(

191 scale: float = None,

192 initial_state: torch.Tensor = None,

193 inplace_final_state: bool = True,

194 - cu_seqlens: torch.LongTensor | None = None,

194 + cu_seqlens: torch.Tensor | None = None,

195 ssm_state_indices: torch.Tensor | None = None,

196 num_accepted_tokens: torch.Tensor | None = None,

197 use_qk_l2norm_in_kernel: bool = False,

vllm/model_executor/layers/fla/ops/index.py

+3 -7

@@ -15,14 +15,12 @@ from .utils import tensor_cache

17 @tensor_cache

18 -def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor:

18 +def prepare_lens(cu_seqlens: torch.Tensor) -> torch.Tensor:

19 return cu_seqlens[1:] - cu_seqlens[:-1]

22 @tensor_cache

23 -def prepare_chunk_indices(

24 - cu_seqlens: torch.LongTensor, chunk_size: int

25 -) -> torch.LongTensor:

23 +def prepare_chunk_indices(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor:

24 indices = torch.cat(

25 [

26 torch.arange(n)

@@ -33,9 +31,7 @@ def prepare_chunk_indices(

33 @tensor_cache

36 -def prepare_chunk_offsets(

37 - cu_seqlens: torch.LongTensor, chunk_size: int

38 -) -> torch.LongTensor:

34 +def prepare_chunk_offsets(cu_seqlens: torch.Tensor, chunk_size: int) -> torch.Tensor:

35 return torch.cat(

36 [cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]

37 ).cumsum(-1)

vllm/model_executor/layers/fla/ops/kda.py

+8 -8

@@ -38,7 +38,7 @@ def fused_recurrent_kda_fwd(

38 scale: float,

39 initial_state: torch.Tensor,

40 inplace_final_state: bool = True,

41 - cu_seqlens: torch.LongTensor | None = None,

41 + cu_seqlens: torch.Tensor | None = None,

42 ssm_state_indices: torch.Tensor | None = None,

43 num_accepted_tokens: torch.Tensor | None = None,

44 use_qk_l2norm_in_kernel: bool = False,

@@ -116,7 +116,7 @@ def fused_recurrent_kda(

116 initial_state: torch.Tensor = None,

117 inplace_final_state: bool = True,

118 use_qk_l2norm_in_kernel: bool = True,

119 - cu_seqlens: torch.LongTensor | None = None,

119 + cu_seqlens: torch.Tensor | None = None,

120 ssm_state_indices: torch.LongTensor | None = None,

121 **kwargs,

122 ) -> tuple[torch.Tensor, torch.Tensor]:

@@ -720,7 +720,7 @@ def chunk_kda_scaled_dot_kkt_fwd(

720 gk: torch.Tensor | None = None,

721 beta: torch.Tensor | None = None,

722 scale: float | None = None,

723 - cu_seqlens: torch.LongTensor | None = None,

723 + cu_seqlens: torch.Tensor | None = None,

724 chunk_size: int = 64,

725 output_dtype: torch.dtype = torch.float32,

726 ) -> tuple[torch.Tensor, torch.Tensor]:

@@ -734,7 +734,7 @@ def chunk_kda_scaled_dot_kkt_fwd(

734 The beta tensor of shape `[B, T, H]`.

735 gk (torch.Tensor):

736 The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`.

737 - cu_seqlens (torch.LongTensor):

737 + cu_seqlens (torch.Tensor):

738 The cumulative sequence lengths of the input tensor.

739 Default: None

740 chunk_size (int):

@@ -964,7 +964,7 @@ def recompute_w_u_fwd(

964 A: torch.Tensor,

965 q: torch.Tensor | None = None,

966 gk: torch.Tensor | None = None,

967 - cu_seqlens: torch.LongTensor | None = None,

967 + cu_seqlens: torch.Tensor | None = None,

968 ) -> tuple[torch.Tensor, torch.Tensor]:

969 B, T, H, K, V = *k.shape, v.shape[-1]

970 BT = A.shape[-1]

@@ -1132,7 +1132,7 @@ def chunk_gla_fwd_o_gk(

1132 h: torch.Tensor,

1133 o: torch.Tensor,

1134 scale: float,

1135 - cu_seqlens: torch.LongTensor | None = None,

1135 + cu_seqlens: torch.Tensor | None = None,

1136 chunk_size: int = 64,

1137 ):

1138 B, T, H, K, V = *q.shape, v.shape[-1]

@@ -1176,7 +1176,7 @@ def chunk_kda_fwd(

1176 scale: float,

1177 initial_state: torch.Tensor,

1178 output_final_state: bool,

1179 - cu_seqlens: torch.LongTensor | None = None,

1179 + cu_seqlens: torch.Tensor | None = None,

1180 ):

1181 chunk_size = 64

1182 g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens)

@@ -1236,7 +1236,7 @@ def chunk_kda(

1236 initial_state: torch.Tensor = None,

1237 output_final_state: bool = False,

1238 use_qk_l2norm_in_kernel: bool = False,

1239 - cu_seqlens: torch.LongTensor | None = None,

1239 + cu_seqlens: torch.Tensor | None = None,

1240 **kwargs,

1241 ):

1242 if scale is None:

vllm/model_executor/layers/fla/ops/wy_fast.py

+1 -1

@@ -122,7 +122,7 @@ def recompute_w_u_fwd(

122 beta: torch.Tensor,

123 g_cumsum: torch.Tensor,

124 A: torch.Tensor,

125 - cu_seqlens: torch.LongTensor | None,

125 + cu_seqlens: torch.Tensor | None,

126 ) -> tuple[torch.Tensor, torch.Tensor]:

127 B, T, Hg, K, V = *k.shape, v.shape[-1]

128 H = v.shape[-2]

vllm/model_executor/models/qwen3_next.py

+13 -8

@@ -119,7 +119,7 @@ def fi_chunk_gated_delta_rule(

119 beta: torch.Tensor,

120 initial_state: torch.Tensor,

121 output_final_state: bool,

122 - cu_seqlens: torch.LongTensor | None = None,

122 + cu_seqlens: torch.Tensor | None = None,

123 use_qk_l2norm_in_kernel: bool = True,

124 ):

125 from flashinfer.gdn_prefill import (

@@ -214,7 +214,7 @@ class ChunkGatedDeltaRule(CustomOp):

214 beta: torch.Tensor,

215 initial_state: torch.Tensor,

216 output_final_state: bool,

217 - cu_seqlens: torch.LongTensor | None = None,

217 + cu_seqlens: torch.Tensor | None = None,

218 use_qk_l2norm_in_kernel: bool = True,

219 ):

220 return fi_chunk_gated_delta_rule(

@@ -238,7 +238,7 @@ class ChunkGatedDeltaRule(CustomOp):

238 beta: torch.Tensor,

239 initial_state: torch.Tensor,

240 output_final_state: bool,

241 - cu_seqlens: torch.LongTensor | None = None,

241 + cu_seqlens: torch.Tensor | None = None,

242 use_qk_l2norm_in_kernel: bool = True,

243 ):

244 return fla_chunk_gated_delta_rule(

@@ -755,8 +755,13 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

755 v = torch.randn(

756 1, T, num_v_heads, self.head_v_dim, device=device, dtype=dtype

757 )

758 - g = torch.randn(1, T, num_v_heads, device=device, dtype=dtype)

759 - beta = torch.randn(1, T, num_v_heads, device=device, dtype=dtype)

758 + # NOTE: g and beta must have the same dtypes as during

759 + # inference, so we construct them with the same function

760 + # (fused_gdn_gating). dummy_a and dummy_b are throwaway

761 + # inputs required by that function.

762 + dummy_a = torch.randn(T, num_v_heads, device=device, dtype=dtype)

763 + dummy_b = torch.randn(T, num_v_heads, device=device, dtype=dtype)

764 + g, beta = fused_gdn_gating(self.A_log, dummy_a, dummy_b, self.dt_bias)

765 state = torch.zeros(

766 1,

767 num_v_heads,

@@ -765,7 +770,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

770 device=device,

771 dtype=state_dtype,

772 )

768 - cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.long)

773 + cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.int32)

774

775 try:

776 self.chunk_gated_delta_rule(

@@ -775,7 +780,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

780 g=g,

781 beta=beta,

782 initial_state=state,

778 - output_final_state=False,

783 + output_final_state=True,

784 cu_seqlens=cu_seqlens,

785 use_qk_l2norm_in_kernel=True,

786 )

@@ -795,7 +800,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):

800 self.prefix,

801 )

802 finally:

798 - del q, k, v, g, beta, state, cu_seqlens

803 + del q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens

804

805 torch.accelerator.empty_cache()

806