flashinfer.gdn_fused_decode_step

flashinfer.gdn_fused_decode_step(hidden_states: Tensor, w_ba: Tensor, mixed_qkv: Tensor, conv_weight: Tensor, conv_bias: Tensor, conv_state: Tensor, A_log: Tensor, dt_bias: Tensor, scale: float | None, ssm_state: Tensor, state_indices: Tensor, use_qk_l2norm: bool = True, out: Tensor | None = None) → Tuple[Tensor, Tensor, Tensor]

Fused single-token GDN decode step over paged conv/ssm state pools.

Folds the per-layer decode chain (b/a projection GEMV, causal conv1d update, q/k/v split, gated delta-rule decode with qk-L2-norm) into one operation. Both state pools are updated in place and returned.

Parameters:

  • hidden_states (torch.Tensor) – Layer input of shape [B, hidden], bfloat16.

  • w_ba (torch.Tensor) – Fused b/a projection weight of shape [hidden, 2*HV], bfloat16 (columns [:HV] produce the beta-gate input, [HV:] the decay input).

  • mixed_qkv (torch.Tensor) – Raw (pre-conv) fused q/k/v channels of shape [B, qkv_dim], bfloat16, with qkv_dim = (2*H_q + HV) * D.

  • conv_weight (torch.Tensor) – Depthwise conv weight of shape [qkv_dim, width], bfloat16.

  • conv_bias (torch.Tensor) – Conv bias of shape [qkv_dim], bfloat16.

  • conv_state (torch.Tensor) – Paged conv-state pool as a logical [P, qkv_dim, width-1] view holding the last width-1 raw channel inputs, bfloat16. Updated in place. Two physical pool layouts are supported: an SD pool ((width-1, qkv_dim) rows, the vLLM default — pass pool.transpose(-1, -2)) or a DS-dense pool ((qkv_dim, width-1) rows, contiguous); the page stride may be padded.

  • A_log (torch.Tensor) – Log decay parameter of shape [HV], float32.

  • dt_bias (torch.Tensor) – Decay bias of shape [HV], bfloat16.

  • scale (float, optional) – Query scale. None and 0.0 both select the default 1/sqrt(D): a zero scale would make the whole attention output zero, so it is treated as “unset” rather than honoured (frameworks that keep the scale in a config default it to 0). Pass an explicit non-zero value to override.

  • ssm_state (torch.Tensor) – Paged fp32 recurrent-state pool of shape [P, HV, V, K] (V-major / K-last), row stride may be padded (stride(0) >= HV*V*K). Updated in place.

  • state_indices (torch.Tensor) –

    Per-batch pool slot indices of shape [B], int32. Both pools are indexed with the same value, so each batch entry reads and writes the same slot.

    Padding / inactive rows: a negative index (vLLM’s PAD_SLOT_ID is -1) marks a batch entry that owns no pool slot — the rows a CUDA-graph replay carries between the live request count and the captured batch size. Such an entry is skipped entirely: neither conv_state nor ssm_state is read or written for it, and its output row is written as zero. This is the same contract the float32 path of gated_delta_rule_decode_pretranspose() documents, so the two FlashInfer GDN decode entry points treat padding identically. The check is in-kernel by necessity: inspecting index values on the host costs a device-to-host sync per layer per decode step and is impossible under graph capture, which is the regime that produces padded rows.

    Indices >= P (the pool’s leading extent) are a caller bug, not a padding convention. They are deliberately neither clamped nor skipped — clamping would corrupt a real slot silently and skipping would turn an index-arithmetic error into a silently missing state update — so the resulting access is out of bounds and its effect is undefined. Only negative values mean padding.

  • use_qk_l2norm (bool) – Apply L2 normalization to q and k. Default True.

  • out (torch.Tensor, optional) – Pre-allocated attention output of shape [B, 1, HV, V], bfloat16, dense (contiguous). Written in place and returned when provided (avoids a separate copy into framework-owned output buffers).

Returns:

(output, conv_state, ssm_state) with output of shape [B, 1, HV, V] (bfloat16) and both pools mutated in place.

Return type:

Tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Notes

  • There is no backend selector and no environment gate: the library chooses between its specialized kernels and the composable path from the registry and the device. The choice is observable — gdn_fused_decode_step_supported() answers it before the call, without running anything — but not overridable per call. A framework that does not want this operation simply does not call it.

  • The specialized kernels serve registered traced workload signatures on SM120; on any other device, or for any geometry the registry does not list, this function is exactly the composable torch implementation.

  • A specialized-kernel failure never breaks this op: it warns once, latches that implementation off for the rest of the process, and the call is served by the composable path.

  • CUDA graphs: each specialized implementation compiles lazily on its first eager dispatch of a (batch, scale, conv-state layout) variant; during capture one is recorded only when that variant is already warm, otherwise the (capture-safe) composable path is baked for that shape.