[Perf][GLM-5.3-Flash] Use FlashKDA for KDA chunked prefill (1.7-3.8x faster than the Triton chunk path) (#55737)
Signed-off-by: Jared Wen <jaredwen@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jared Wen committed
Sep 14, 2026 at 16:36 UTC
b443c1cc4e12171be0b0bbb7cc107919bda4634f
1 file changed
+144
-28
vllm/models/glm5next/nvidia/kda.py
+144
-28
@@ -37,6 +37,7 @@ from vllm.platforms import current_platform
37
from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
38
from vllm.transformers_utils.configs.glm5_next import Glm5NextConfig
39
from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
40
+from vllm.v1.worker.workspace import current_workspace_manager
41
42
if current_platform.is_rocm():
43
from vllm.models.glm5next.amd.ops.third_party.kda import (
@@ -123,6 +124,32 @@ def _cast_sigmoid(x: torch.Tensor) -> torch.Tensor:
124
return x.float().sigmoid()
125
126
127
+def _resolve_kda_prefill_backend(
128
+ backend: str, head_dim: int, dtype: torch.dtype, lower_bound: float | None
129
+) -> str:
130
+ """Pick the chunked-prefill kernel: FlashKDA (fused CUDA, ~2-4x faster on
131
+ SM90/SM10x/SM12x for bf16, head_dim 128 and a bounded gate) or the Triton
132
+ ``chunk_kda_with_fused_gate`` path. ``backend`` comes from
133
+ ``additional_config.kda_prefill_backend`` (auto / triton / flashkda)."""
134
+ if backend not in ("auto", "triton", "flashkda"):
135
+ raise ValueError(f"Unsupported KDA prefill backend: {backend}")
136
+ capability = current_platform.get_device_capability()
137
+ supported = (
138
+ current_platform.is_cuda()
139
+ and capability is not None
140
+ and capability.major in (9, 10, 12)
141
+ and head_dim == 128
142
+ and dtype == torch.bfloat16
143
+ and lower_bound is not None
144
+ )
145
+ if backend == "flashkda" and not supported:
146
+ raise RuntimeError(
147
+ "FlashKDA requires CUDA SM90/SM10x/SM12x, bfloat16, head_dim=128 "
148
+ "and a bounded KDA gate."
149
+ )
150
+ return "flashkda" if supported and backend != "triton" else "triton"
151
+
152
+
153
class Glm5NextLinearAttention(GatedDeltaNetAttention):
154
head_dim: int
155
num_heads: int
@@ -291,6 +318,85 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):
318
# every _forward call (it reads an env-derived flag each time).
319
self._conv_state_dim_first = is_conv_state_dim_first()
320
321
+ additional_config = vllm_config.additional_config
322
+ self.kda_prefill_backend = _resolve_kda_prefill_backend(
323
+ additional_config.get("kda_prefill_backend", "auto")
324
+ if isinstance(additional_config, dict)
325
+ else "auto",
326
+ self.head_dim,
327
+ vllm_config.model_config.dtype,
328
+ self.kda_lower_bound,
329
+ )
330
+ self._flashkda_buffer_specs: (
331
+ tuple[tuple[tuple[int, ...], torch.dtype], ...] | None
332
+ ) = None
333
+ if self.kda_prefill_backend == "flashkda":
334
+ import vllm._flashkda_C # noqa: F401
335
+
336
+ max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
337
+ max_seqs = vllm_config.scheduler_config.max_num_seqs
338
+ workspace_size = torch.ops._flashkda_C.get_workspace_size(
339
+ max_tokens, self.local_num_heads, max_seqs
340
+ )
341
+ self._flashkda_buffer_specs = (
342
+ (
343
+ (max_seqs, self.local_num_heads, self.head_dim, self.head_dim),
344
+ self.get_state_dtype()[1],
345
+ ),
346
+ ((workspace_size,), torch.uint8),
347
+ # Output buffer for steps that also carry spec-decode tokens:
348
+ # the non-spec tokens are then scattered by non_spec_token_indx.
349
+ (
350
+ (1, max_tokens, self.local_num_heads, self.head_dim),
351
+ vllm_config.model_config.dtype,
352
+ ),
353
+ )
354
+
355
+ def _flashkda_prefill(
356
+ self,
357
+ q: torch.Tensor,
358
+ k: torch.Tensor,
359
+ v: torch.Tensor,
360
+ g: torch.Tensor,
361
+ beta: torch.Tensor,
362
+ initial_state: torch.Tensor,
363
+ cu_seqlens: torch.Tensor,
364
+ out: torch.Tensor | None,
365
+ ) -> tuple[torch.Tensor, torch.Tensor]:
366
+ """Fused KDA chunked prefill (FlashKDA). Takes the raw gate logits ``g``
367
+ and raw ``beta`` logits, l2-normalizes q/k in-kernel and applies the
368
+ bounded gate ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``,
369
+ matching ``chunk_kda_with_fused_gate(..., safe_gate=True)``. Writes the
370
+ attention output into ``out`` (a workspace buffer when ``None``) and
371
+ returns ``(out, final_state)``."""
372
+ assert self._flashkda_buffer_specs is not None
373
+ final_state, workspace, workspace_out = (
374
+ current_workspace_manager().get_simultaneous(*self._flashkda_buffer_specs)
375
+ )
376
+ final_state = final_state[: initial_state.shape[0]]
377
+ if out is None:
378
+ out = workspace_out[:, : q.shape[1]]
379
+ # FlashKDA hardcodes dense q/k/v/g strides; beta may be row-strided.
380
+ torch.ops._flashkda_C.fwd(
381
+ q.contiguous(),
382
+ k.contiguous(),
383
+ v.contiguous(),
384
+ g.contiguous(),
385
+ beta,
386
+ self.head_dim**-0.5,
387
+ out,
388
+ workspace,
389
+ self.A_log.view(-1),
390
+ self.dt_bias.view(-1, self.head_dim),
391
+ self.kda_lower_bound,
392
+ initial_state.contiguous(),
393
+ final_state,
394
+ cu_seqlens.contiguous(),
395
+ None,
396
+ None,
397
+ )
398
+ return out, final_state
399
+
400
def forward(
401
self,
402
hidden_states: torch.Tensor,
@@ -542,26 +648,42 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):
648
initial_state = gather_initial_states(
649
recurrent_state, non_spec_state_indices_tensor, has_initial_state
650
)
545
- (
546
- core_attn_out_non_spec,
547
- last_recurrent_state,
548
- ) = chunk_kda_with_fused_gate(
549
- q=_rearr(q_ns),
550
- k=_rearr(k_ns),
551
- v=_rearr(v_ns),
552
- raw_g=g1_ns,
553
- # Chunk path wants the pre-sigmoided fp32 beta (its kernels
554
- # don't sigmoid); beta_ns is raw bf16 from forward.
555
- beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0),
556
- A_log=self.A_log,
557
- g_bias=self.dt_bias,
558
- initial_state=initial_state,
559
- output_final_state=True,
560
- use_qk_l2norm_in_kernel=True,
561
- cu_seqlens=non_spec_query_start_loc,
562
- safe_gate=safe_gate,
563
- lower_bound=lower_bound,
564
- )
651
+ if self.kda_prefill_backend == "flashkda":
652
+ # Non-spec step: write straight into the layer output buffer
653
+ # (dense token order, no merge copy). Step with spec-decode
654
+ # tokens: write to the workspace buffer and scatter below.
655
+ ns_out = None if use_spec else core_attn_out[:, :num_actual_tokens]
656
+ core_attn_out_non_spec, last_recurrent_state = self._flashkda_prefill(
657
+ q=_rearr(q_ns),
658
+ k=_rearr(k_ns),
659
+ v=_rearr(v_ns),
660
+ g=g1_ns,
661
+ beta=beta_ns,
662
+ initial_state=initial_state,
663
+ cu_seqlens=non_spec_query_start_loc,
664
+ out=ns_out,
665
+ )
666
+ else:
667
+ (
668
+ core_attn_out_non_spec,
669
+ last_recurrent_state,
670
+ ) = chunk_kda_with_fused_gate(
671
+ q=_rearr(q_ns),
672
+ k=_rearr(k_ns),
673
+ v=_rearr(v_ns),
674
+ raw_g=g1_ns,
675
+ # Chunk path wants the pre-sigmoided fp32 beta (its
676
+ # kernels don't sigmoid); beta_ns is raw bf16 from forward.
677
+ beta=_cast_sigmoid(beta_ns.squeeze(0)).unsqueeze(0),
678
+ A_log=self.A_log,
679
+ g_bias=self.dt_bias,
680
+ initial_state=initial_state,
681
+ output_final_state=True,
682
+ use_qk_l2norm_in_kernel=True,
683
+ cu_seqlens=non_spec_query_start_loc,
684
+ safe_gate=safe_gate,
685
+ lower_bound=lower_bound,
686
+ )
687
# Init cache
688
scatter_states(
689
recurrent_state,
@@ -600,14 +722,8 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):
722
# --- merge spec / non-spec outputs back into token order ---
723
if use_spec and core_attn_out_non_spec is not None:
724
assert core_attn_out_spec is not None
603
- merged = torch.empty(
604
- (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
605
- dtype=core_attn_out_non_spec.dtype,
606
- device=core_attn_out_non_spec.device,
607
- )
608
- merged.index_copy_(1, spec_token_indx, core_attn_out_spec)
609
- merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
610
- core_attn_out[0, :num_actual_tokens] = merged.squeeze(0)
725
+ core_attn_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
726
+ core_attn_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
727
elif use_spec:
728
assert core_attn_out_spec is not None
729
if spec_out is None: