Skip to content

perf(gdn): make GDN kernels compilation batch-size agnostic (support dynamic batch shapes for vLLM integration) #3649

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bkryu merged 6 commits into from
Jun 22, 2026
Merged

perf(gdn): make GDN kernels compilation batch-size agnostic (support dynamic batch shapes for vLLM integration) #3649

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
2527568
perf(gdn): make BF16 decode kernels batch-size agnostic
kahyunnam Jun 16, 2026
4995f9f
perf(gdn): make pretranspose/nontranspose/MTP decode kernels batch-si…
kahyunnam Jun 16, 2026
12785fc
perf(gdn): make BF16 decode pool-size agnostic for contiguous pools
kahyunnam Jun 16, 2026
4c06561
perf(gdn): guard caching-off dummy intermediate_states marking in MTP
kahyunnam Jun 16, 2026
228d811
style(gdn): trim comment bloat in decode dynamic-batch changes
kahyunnam Jun 16, 2026
3322c7f
fix(gdn): use full stride tuple in BF16 non-contiguous pool cache key
kahyunnam Jun 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 74 additions & 74 deletions flashinfer/gdn_kernels/gdn_decode_bf16_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,15 @@
from flashinfer.cute_dsl.fp4_common import get_sm_version
from flashinfer.cute_dsl.utils import get_num_sm


def _mark_batch_dynamic(torch_t: torch.Tensor, *, assumed_align: int = 32):
# explicit stride_order: auto-deduction is ambiguous at B=1/T=1.
stride_order = tuple(sorted(range(torch_t.dim()), key=lambda d: -torch_t.stride(d)))
return from_dlpack(
torch_t, assumed_align=assumed_align, enable_tvm_ffi=True
).mark_compact_shape_dynamic(mode=0, stride_order=stride_order)


# ==============================================================================
# FMA WRAPPER FUNCTIONS (SM90 Compatibility)
# ==============================================================================
Expand Down Expand Up @@ -129,7 +138,6 @@ def gdn_decode_bf16state_mtp_ilp4_kernel(
softplus_threshold: cutlass.Constexpr[float],
scale: cutlass.Constexpr[float],
HV: cutlass.Constexpr[int],
B: cutlass.Constexpr[int],
T: cutlass.Constexpr[int],
H: cutlass.Constexpr[int],
K: cutlass.Constexpr[int],
Expand Down Expand Up @@ -749,7 +757,6 @@ def gdn_wide_vec_kernel(
softplus_threshold: cutlass.Constexpr[float],
scale: cutlass.Constexpr[float],
HV: cutlass.Constexpr[int],
B: cutlass.Constexpr[int],
T: cutlass.Constexpr[int],
H: cutlass.Constexpr[int],
K: cutlass.Constexpr[int],
Expand Down Expand Up @@ -1248,7 +1255,6 @@ def run_gdn_decode_bf16state_mtp_ilp4(
softplus_threshold: cutlass.Constexpr[float],
scale: cutlass.Constexpr[float],
HV: cutlass.Constexpr[int],
B: cutlass.Constexpr[int],
T: cutlass.Constexpr[int],
H: cutlass.Constexpr[int],
K: cutlass.Constexpr[int],
Expand All @@ -1273,6 +1279,7 @@ def run_gdn_decode_bf16state_mtp_ilp4(
)

num_v_tiles = cute.ceil_div(v_dim, tile_v)
B = cute.size(q.shape[0])
grid_size = B * HV * num_v_tiles

smem_bytes = 128
Expand All @@ -1299,7 +1306,6 @@ def run_gdn_decode_bf16state_mtp_ilp4(
softplus_threshold,
scale,
HV,
B,
T,
H,
K,
Expand Down Expand Up @@ -1340,7 +1346,6 @@ def _run_wide_vec(
softplus_threshold: cutlass.Constexpr[float],
scale: cutlass.Constexpr[float],
HV: cutlass.Constexpr[int],
B: cutlass.Constexpr[int],
T: cutlass.Constexpr[int],
H: cutlass.Constexpr[int],
K: cutlass.Constexpr[int],
Expand All @@ -1354,6 +1359,7 @@ def _run_wide_vec(
stream: cuda.CUstream,
):
num_v_tiles: cutlass.Constexpr[int] = V // tile_v
B = cute.size(q.shape[0])
grid_size = B * HV * num_v_tiles
smem_bytes = (
4 * T * (K + 8) # sQ FP32
Expand All @@ -1378,7 +1384,6 @@ def _run_wide_vec(
softplus_threshold,
scale,
HV,
B,
T,
H,
K,
Expand Down Expand Up @@ -1721,19 +1726,23 @@ def gated_delta_rule_mtp_wide_vec(
output_state_indices is None or output_state_indices is initial_state_indices
)

# Slot stride may be larger than HV*V*K (vLLM packs conv+ssm on the same
# page). Include in the cache key so padded vs tight pools each get their
# own compiled kernel — cute.compile bakes the stride into the cubin.
pool_slot_stride = int(initial_state_source.stride(0))
# Contiguous pool -> sentinel keys + slot dim marked dynamic (pool-size
# agnostic); padded/strided pool keeps real pool_size/stride in the key.
contiguous_pool = initial_state_source.is_contiguous()
if contiguous_pool:
pool_size_key = -1
pool_slot_stride: tuple[int, ...] = (-1,)
else:
pool_size_key = pool_size
pool_slot_stride = tuple(int(s) for s in initial_state_source.stride())
cache_key = (
"v3_mtp_bf16_tiled",
B_val,
"v3_mtp_bf16_tiled_dynB",
T_val,
H_val,
HV_val,
K_val,
V_val,
pool_size,
pool_size_key,
pool_slot_stride,
tile_v,
effective_disable_final,
Expand All @@ -1745,36 +1754,33 @@ def gated_delta_rule_mtp_wide_vec(
use_packed_fma,
same_pool,
)
if cache_key not in _compiled_kernels_wide_vec:
default_indices = torch.arange(B_val, dtype=torch.int32, device=q.device)
default_output = torch.empty(

if initial_state_indices is None:
initial_state_indices = torch.arange(B_val, dtype=torch.int32, device=q.device)
if output is None:
output = torch.empty(
B_val, T_val, HV_val, V_val, device=q.device, dtype=q.dtype
)

if initial_state_indices is None:
initial_state_indices = default_indices
if output is None:
output = default_output

# Compile-time indices template: any [B] int32 on the right device.
# Both read and write index tensors share the same shape/dtype so
# the same dlpack handle works for both slots.
h_ = from_dlpack(h0_source, assumed_align=32, enable_tvm_ffi=True)
if cache_key not in _compiled_kernels_wide_vec:
if contiguous_pool:
h_ = _mark_batch_dynamic(h0_source)
else:
h_ = from_dlpack(h0_source, assumed_align=32, enable_tvm_ffi=True)
inter_ = from_dlpack(intermediate_states, assumed_align=32, enable_tvm_ffi=True)
q_ = from_dlpack(q, assumed_align=32, enable_tvm_ffi=True)
k_ = from_dlpack(k, assumed_align=32, enable_tvm_ffi=True)
v_ = from_dlpack(v, assumed_align=32, enable_tvm_ffi=True)
a_ = from_dlpack(a, assumed_align=32, enable_tvm_ffi=True)
b_ = from_dlpack(b, assumed_align=32, enable_tvm_ffi=True)
if cache_intermediate_states:
# Dummy [1,1,1] tensor (caching off) has no unique stride-1 dim.
inter_ = _mark_batch_dynamic(intermediate_states)
q_ = _mark_batch_dynamic(q)
k_ = _mark_batch_dynamic(k)
v_ = _mark_batch_dynamic(v)
a_ = _mark_batch_dynamic(a)
b_ = _mark_batch_dynamic(b)
A_log_ = from_dlpack(A_log, assumed_align=32, enable_tvm_ffi=True)
dt_bias_ = from_dlpack(dt_bias, assumed_align=32, enable_tvm_ffi=True)
o_ = from_dlpack(output, assumed_align=32, enable_tvm_ffi=True)
h0_idx_ = from_dlpack(
initial_state_indices, assumed_align=32, enable_tvm_ffi=True
)
h0_out_idx_ = from_dlpack(
initial_state_indices, assumed_align=32, enable_tvm_ffi=True
)
o_ = _mark_batch_dynamic(output)
h0_idx_ = _mark_batch_dynamic(initial_state_indices)
h0_out_idx_ = _mark_batch_dynamic(initial_state_indices)

_compiled_kernels_wide_vec[cache_key] = {
"compiled": cute.compile(
Expand All @@ -1795,7 +1801,6 @@ def gated_delta_rule_mtp_wide_vec(
softplus_threshold,
scale,
HV_val,
B_val,
T_val,
H_val,
K_val,
Expand All @@ -1809,19 +1814,13 @@ def gated_delta_rule_mtp_wide_vec(
stream,
options="--enable-tvm-ffi --generate-line-info --opt-level 3",
),
"default_indices": default_indices,
"output": default_output,
}

cache = _compiled_kernels_wide_vec[cache_key]
if initial_state_indices is None:
initial_state_indices = cache["default_indices"]
if output_state_indices is None:
# Single-pool: read==write. Reuse the same indices tensor — no extra
# allocation, kernel still produces the same address for both slots.
output_state_indices = initial_state_indices
if output is None:
output = cache["output"]

cache["compiled"](
h0_source,
Expand Down Expand Up @@ -1987,18 +1986,23 @@ def gated_delta_rule_mtp(
output_state_indices is None or output_state_indices is initial_state_indices
)

# Slot stride may be larger than HV*V*K (vLLM's packed conv+ssm page).
# Include in cache key — cute.compile bakes the stride into the cubin.
pool_slot_stride = int(initial_state_source.stride(0))
# Contiguous pool -> sentinel keys + slot dim marked dynamic (pool-size
# agnostic); padded/strided pool keeps real pool_size/stride in the key.
contiguous_pool = initial_state_source.is_contiguous()
if contiguous_pool:
pool_size_key = -1
pool_slot_stride: tuple[int, ...] = (-1,)
else:
pool_size_key = pool_size
pool_slot_stride = tuple(int(s) for s in initial_state_source.stride())
cache_key = (
"mtp_bf16",
B,
"mtp_bf16_dynB",
T,
H,
HV,
K,
V,
pool_size,
pool_size_key,
pool_slot_stride,
tile_v,
ilp_rows,
Expand All @@ -2011,29 +2015,30 @@ def gated_delta_rule_mtp(
use_packed_fma,
same_pool,
)

if initial_state_indices is None:
initial_state_indices = torch.arange(B, dtype=torch.int32, device=q.device)
if output is None:
output = torch.empty(B, T, HV, V, device=q.device, dtype=q.dtype)

if cache_key not in _compiled_kernels_mtp:
# First call for this shape: allocate default indices/output and do
# dlpack conversions once for compilation. Steady-state calls pass
# torch tensors straight to the compiled callable (tvm-ffi accepts
# either) and reuse these cached defaults when the caller doesn't
# provide their own.
default_indices = torch.arange(B, dtype=torch.int32, device=q.device)
default_output = torch.empty(B, T, HV, V, device=q.device, dtype=q.dtype)

h_ = from_dlpack(h0_source, assumed_align=32, enable_tvm_ffi=True)
if contiguous_pool:
h_ = _mark_batch_dynamic(h0_source)
else:
h_ = from_dlpack(h0_source, assumed_align=32, enable_tvm_ffi=True)
inter_ = from_dlpack(intermediate_states, assumed_align=32, enable_tvm_ffi=True)
q_ = from_dlpack(q, assumed_align=32, enable_tvm_ffi=True)
k_ = from_dlpack(k, assumed_align=32, enable_tvm_ffi=True)
v_ = from_dlpack(v, assumed_align=32, enable_tvm_ffi=True)
a_ = from_dlpack(a, assumed_align=32, enable_tvm_ffi=True)
b_ = from_dlpack(b, assumed_align=32, enable_tvm_ffi=True)
if cache_intermediate_states:
inter_ = _mark_batch_dynamic(intermediate_states)
q_ = _mark_batch_dynamic(q)
k_ = _mark_batch_dynamic(k)
v_ = _mark_batch_dynamic(v)
a_ = _mark_batch_dynamic(a)
b_ = _mark_batch_dynamic(b)
A_log_ = from_dlpack(A_log, assumed_align=32, enable_tvm_ffi=True)
dt_bias_ = from_dlpack(dt_bias, assumed_align=32, enable_tvm_ffi=True)
o_ = from_dlpack(default_output, assumed_align=32, enable_tvm_ffi=True)
h0_idx_ = from_dlpack(default_indices, assumed_align=32, enable_tvm_ffi=True)
h0_out_idx_ = from_dlpack(
default_indices, assumed_align=32, enable_tvm_ffi=True
)
o_ = _mark_batch_dynamic(output)
h0_idx_ = _mark_batch_dynamic(initial_state_indices)
h0_out_idx_ = _mark_batch_dynamic(initial_state_indices)

_compiled_kernels_mtp[cache_key] = {
"compiled": cute.compile(
Expand All @@ -2054,7 +2059,6 @@ def gated_delta_rule_mtp(
softplus_threshold,
scale,
HV,
B,
T,
H,
K,
Expand All @@ -2068,16 +2072,12 @@ def gated_delta_rule_mtp(
stream,
options="--enable-tvm-ffi --generate-line-info --opt-level 3",
),
"default_indices": default_indices,
"output": default_output,
}

cache = _compiled_kernels_mtp[cache_key]

if output_state_indices is None:
output_state_indices = initial_state_indices
if output is None:
output = cache["output"]

cache["compiled"](
h0_source,
Expand Down
Loading
Loading