From 252723cd074d1afd7f0cb418d4d39e19c19069ef Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Sun, 6 Sep 2026 06:43:02 +0000 Subject: [PATCH 1/2] [Perf][GLM-5.3-Flash] Use FlashKDA for KDA chunked prefill FlashKDA (already built as vllm._flashkda_C for Kimi-K3) implements the same bounded-gate KDA recurrence as chunk_kda_with_fused_gate (lower_bound * sigmoid(exp(A_log) * (g + dt_bias)), in-kernel q/k l2norm, raw beta logits). On GB300 it replaces ~15 Triton kernels with 4 and is 1.7-3.8x faster (T=2048x8: 183->49 us, 8192x4: 555->146 us, 16384x1: 1371->787 us per layer). Select it automatically when supported (SM90/SM10x/SM12x, bf16, head_dim 128, bounded gate); additional_config.kda_prefill_backend = triton keeps the old path. The output is written straight into the layer buffer and the pre-sigmoid beta cast is no longer needed on this path. Co-authored-by: Claude Fable 5.1 Signed-off-by: Jared Wen --- vllm/models/glm5next/nvidia/kda.py | 153 +++++++++++++++++++++++++---- 1 file changed, 133 insertions(+), 20 deletions(-) diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index e6ba6ac5c659..5cc5fdafb87e 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/vllm/models/glm5next/nvidia/kda.py @@ -37,6 +37,7 @@ from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata +from vllm.v1.worker.workspace import current_workspace_manager if current_platform.is_rocm(): from vllm.models.glm5next.amd.ops.third_party.kda import ( @@ -123,6 +124,32 @@ def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor: return x.float().sigmoid() +def _resolve_kda_prefill_backend( + backend: str, head_dim: int, dtype: torch.dtype, lower_bound: float | None +) -> str: + """Pick the chunked-prefill kernel: FlashKDA (fused CUDA, ~2-4x faster on + SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate) or the Triton + ``chunk_kda_with_fused_gate`` path. ``backend`` comes from + ``additional_config.kda_prefill_backend`` (auto / triton / flashkda).""" + if backend not in ("auto", "triton", "flashkda"): + raise ValueError(f"Unsupported KDA prefill backend: {backend}") + capability = current_platform.get_device_capability() + supported = ( + current_platform.is_cuda() + and capability is not None + and capability.major in (9, 10, 12) + and head_dim == 128 + and dtype == torch.bfloat16 + and lower_bound is not None + ) + if backend == "flashkda" and not supported: + raise RuntimeError( + "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, head_dim=128 " + "and a bounded KDA gate." + ) + return "flashkda" if supported and backend != "triton" else "triton" + + class Glm5NextLinearAttention(GatedDeltaNetAttention): head_dim: int num_heads: int @@ -282,6 +309,76 @@ def _a_log_weight_loader(param, loaded_weight): # every _forward call (it reads an env-derived flag each time). self._conv_state_dim_first = is_conv_state_dim_first() + additional_config = vllm_config.additional_config + self.kda_prefill_backend = _resolve_kda_prefill_backend( + additional_config.get("kda_prefill_backend", "auto") + if isinstance(additional_config, dict) + else "auto", + self.head_dim, + vllm_config.model_config.dtype, + self.kda_lower_bound, + ) + self._flashkda_buffer_specs: ( + tuple[tuple[tuple[int, ...], torch.dtype], ...] | None + ) = None + if self.kda_prefill_backend == "flashkda": + import vllm._flashkda_C # noqa: F401 + + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + max_seqs = vllm_config.scheduler_config.max_num_seqs + workspace_size = torch.ops._flashkda_C.get_workspace_size( + max_tokens, self.local_num_heads, max_seqs + ) + self._flashkda_buffer_specs = ( + ( + (max_seqs, self.local_num_heads, self.head_dim, self.head_dim), + self.get_state_dtype()[1], + ), + ((workspace_size,), torch.uint8), + ) + + def _flashkda_prefill( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + out: torch.Tensor, + ) -> torch.Tensor: + """Fused KDA chunked prefill (FlashKDA). Takes the raw gate logits ``g`` + and raw ``beta`` logits, l2-normalizes q/k in-kernel and applies the + bounded gate ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, + matching ``chunk_kda_with_fused_gate(..., safe_gate=True)``. Writes the + attention output into ``out`` and returns the final states.""" + assert self._flashkda_buffer_specs is not None + final_state, workspace = current_workspace_manager().get_simultaneous( + *self._flashkda_buffer_specs + ) + final_state = final_state[: initial_state.shape[0]] + # FlashKDA hardcodes dense q/k/v/g strides; beta may be row-strided. + torch.ops._flashkda_C.fwd( + q.contiguous(), + k.contiguous(), + v.contiguous(), + g.contiguous(), + beta, + self.head_dim**-0.5, + out, + workspace, + self.A_log.view(-1), + self.dt_bias.view(-1, self.head_dim), + self.kda_lower_bound, + initial_state.contiguous(), + final_state, + cu_seqlens.contiguous(), + None, + None, + ) + return final_state + def forward( self, hidden_states: torch.Tensor, @@ -533,26 +630,42 @@ def _rearr(x): initial_state = gather_initial_states( recurrent_state, non_spec_state_indices_tensor, has_initial_state ) - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = chunk_kda_with_fused_gate( - q=_rearr(q_ns), - k=_rearr(k_ns), - v=_rearr(v_ns), - raw_g=g1_ns, - # Chunk path wants the pre-sigmoided fp32 beta (its kernels - # don't sigmoid); beta_ns is raw bf16 from forward. - beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0), - A_log=self.A_log, - g_bias=self.dt_bias, - initial_state=initial_state, - output_final_state=True, - use_qk_l2norm_in_kernel=True, - cu_seqlens=non_spec_query_start_loc, - safe_gate=safe_gate, - lower_bound=lower_bound, - ) + if self.kda_prefill_backend == "flashkda" and not use_spec: + # Write straight into the layer output buffer (dense token + # order in a non-spec step); no merge copy below. + ns_out = core_attn_out[:, :num_actual_tokens] + last_recurrent_state = self._flashkda_prefill( + q=_rearr(q_ns), + k=_rearr(k_ns), + v=_rearr(v_ns), + g=g1_ns, + beta=beta_ns, + initial_state=initial_state, + cu_seqlens=non_spec_query_start_loc, + out=ns_out, + ) + core_attn_out_non_spec = ns_out + else: + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = chunk_kda_with_fused_gate( + q=_rearr(q_ns), + k=_rearr(k_ns), + v=_rearr(v_ns), + raw_g=g1_ns, + # Chunk path wants the pre-sigmoided fp32 beta (its + # kernels don't sigmoid); beta_ns is raw bf16 from forward. + beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0), + A_log=self.A_log, + g_bias=self.dt_bias, + initial_state=initial_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=non_spec_query_start_loc, + safe_gate=safe_gate, + lower_bound=lower_bound, + ) # Init cache scatter_states( recurrent_state, From 674dd959c04dec16e3a0cb7cb59c5eebde87cdb5 Mon Sep 17 00:00:00 2001 From: Jared Wen Date: Mon, 7 Sep 2026 14:19:51 +0000 Subject: [PATCH 2/2] [Perf][GLM-5.3-Flash] Use FlashKDA for the prefill segment of spec-decode steps too The FlashKDA path was gated on `not use_spec` so it could write straight into the layer output buffer; a step that also carried spec-decode tokens fell back to the Triton chunk path for its whole prefill segment, which with MTP enabled is almost every step that has a prefill. Follow the Kimi-K3 KDA layer instead: in a spec step FlashKDA writes to a workspace buffer and the non-spec tokens are scattered by non_spec_token_indx, alongside the spec tokens (which keep the recurrent kernel with num_accepted_tokens rollback). The merge now index_copy_s directly into core_attn_out instead of going through a temporary tensor. MTP k=1, TP4 on 4x GB300, same build, Triton -> FlashKDA: 8x2048 TTFT 464 -> 423 ms, 2x8192 390 -> 360 ms, 32k/256 c=16 260 -> 284 tok/s, 1k/512 c=64 2003 -> 2130 tok/s; gsm8k 93.33 vs 93.56 (+-0.7). Co-authored-by: Claude Fable 5.1 Signed-off-by: Jared Wen --- vllm/models/glm5next/nvidia/kda.py | 43 ++++++++++++++++-------------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/vllm/models/glm5next/nvidia/kda.py b/vllm/models/glm5next/nvidia/kda.py index 5cc5fdafb87e..45e1adf75066 100644 --- a/vllm/models/glm5next/nvidia/kda.py +++ b/vllm/models/glm5next/nvidia/kda.py @@ -335,6 +335,12 @@ def _a_log_weight_loader(param, loaded_weight): self.get_state_dtype()[1], ), ((workspace_size,), torch.uint8), + # Output buffer for steps that also carry spec-decode tokens: + # the non-spec tokens are then scattered by non_spec_token_indx. + ( + (1, max_tokens, self.local_num_heads, self.head_dim), + vllm_config.model_config.dtype, + ), ) def _flashkda_prefill( @@ -346,18 +352,21 @@ def _flashkda_prefill( beta: torch.Tensor, initial_state: torch.Tensor, cu_seqlens: torch.Tensor, - out: torch.Tensor, - ) -> torch.Tensor: + out: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: """Fused KDA chunked prefill (FlashKDA). Takes the raw gate logits ``g`` and raw ``beta`` logits, l2-normalizes q/k in-kernel and applies the bounded gate ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``, matching ``chunk_kda_with_fused_gate(..., safe_gate=True)``. Writes the - attention output into ``out`` and returns the final states.""" + attention output into ``out`` (a workspace buffer when ``None``) and + returns ``(out, final_state)``.""" assert self._flashkda_buffer_specs is not None - final_state, workspace = current_workspace_manager().get_simultaneous( - *self._flashkda_buffer_specs + final_state, workspace, workspace_out = ( + current_workspace_manager().get_simultaneous(*self._flashkda_buffer_specs) ) final_state = final_state[: initial_state.shape[0]] + if out is None: + out = workspace_out[:, : q.shape[1]] # FlashKDA hardcodes dense q/k/v/g strides; beta may be row-strided. torch.ops._flashkda_C.fwd( q.contiguous(), @@ -377,7 +386,7 @@ def _flashkda_prefill( None, None, ) - return final_state + return out, final_state def forward( self, @@ -630,11 +639,12 @@ def _rearr(x): initial_state = gather_initial_states( recurrent_state, non_spec_state_indices_tensor, has_initial_state ) - if self.kda_prefill_backend == "flashkda" and not use_spec: - # Write straight into the layer output buffer (dense token - # order in a non-spec step); no merge copy below. - ns_out = core_attn_out[:, :num_actual_tokens] - last_recurrent_state = self._flashkda_prefill( + if self.kda_prefill_backend == "flashkda": + # Non-spec step: write straight into the layer output buffer + # (dense token order, no merge copy). Step with spec-decode + # tokens: write to the workspace buffer and scatter below. + ns_out = None if use_spec else core_attn_out[:, :num_actual_tokens] + core_attn_out_non_spec, last_recurrent_state = self._flashkda_prefill( q=_rearr(q_ns), k=_rearr(k_ns), v=_rearr(v_ns), @@ -644,7 +654,6 @@ def _rearr(x): cu_seqlens=non_spec_query_start_loc, out=ns_out, ) - core_attn_out_non_spec = ns_out else: ( core_attn_out_non_spec, @@ -704,14 +713,8 @@ def _rearr(x): # --- merge spec / non-spec outputs back into token order --- if use_spec and core_attn_out_non_spec is not None: assert core_attn_out_spec is not None - merged = torch.empty( - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), - dtype=core_attn_out_non_spec.dtype, - device=core_attn_out_non_spec.device, - ) - merged.index_copy_(1, spec_token_indx, core_attn_out_spec) - merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) - core_attn_out[0, :num_actual_tokens] = merged.squeeze(0) + core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec) + core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) elif use_spec: assert core_attn_out_spec is not None if spec_out is None: