[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>
This commit is contained in:
Jared Wen
co-authored by
Claude Fable 5.1
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Claude Fable 5.1
mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent
23cfaad497
commit
b443c1cc4e
@@ -37,6 +37,7 @@ from vllm.platforms import current_platform
|
||||
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
|
||||
@@ -291,6 +318,85 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):
|
||||
# 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),
|
||||
# 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(
|
||||
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 | 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`` (a workspace buffer when ``None``) and
|
||||
returns ``(out, final_state)``."""
|
||||
assert self._flashkda_buffer_specs is not None
|
||||
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(),
|
||||
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 out, final_state
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
@@ -542,26 +648,42 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):
|
||||
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":
|
||||
# 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),
|
||||
g=g1_ns,
|
||||
beta=beta_ns,
|
||||
initial_state=initial_state,
|
||||
cu_seqlens=non_spec_query_start_loc,
|
||||
out=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,
|
||||
@@ -600,14 +722,8 @@ class Glm5NextLinearAttention(GatedDeltaNetAttention):
|
||||
# --- 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:
|
||||
|
||||
Reference in New Issue
Block a user