diff --git a/python/sglang/kernels/ops/attention/helion/kda_decode.py b/python/sglang/kernels/ops/attention/helion/kda_decode.py index 8bda2e945..92d7944a3 100644 --- a/python/sglang/kernels/ops/attention/helion/kda_decode.py +++ b/python/sglang/kernels/ops/attention/helion/kda_decode.py @@ -210,6 +210,10 @@ def _select_decode_kernel( return _helion_fused_recurrent_kda_packed_decode +def _is_power_of_two(value: int) -> bool: + return value > 0 and value & (value - 1) == 0 + + def validate_packed_decode_inputs( mixed_qkv: torch.Tensor, a: torch.Tensor, @@ -280,6 +284,11 @@ def validate_packed_decode_inputs( if initial_state.stride(-1) != 1: raise ValueError("`initial_state` must be contiguous in the last dim.") HV, V, K = initial_state.shape[-3:] + if not _is_power_of_two(K) or not _is_power_of_two(V): + raise ValueError( + "Helion KDA decode requires power-of-two key and value head " + f"dimensions (got K={K}, V={V})." + ) if a.shape[1] != HV * K: raise ValueError( f"`a` must have shape [B, HV*K] with HV={HV}, K={K} " diff --git a/python/sglang/kernels/ops/attention/helion/kda_prefill.py b/python/sglang/kernels/ops/attention/helion/kda_prefill.py index 931b370fc..25e0eb3c1 100644 --- a/python/sglang/kernels/ops/attention/helion/kda_prefill.py +++ b/python/sglang/kernels/ops/attention/helion/kda_prefill.py @@ -15,6 +15,7 @@ from sglang.kernels.ops.attention.fla.index import ( prepare_chunk_indices, prepare_chunk_offsets, ) +from sglang.kernels.ops.attention.fla.kda import chunk_kda as triton_chunk_kda CHUNK_SIZE = 64 # Rounded identically to flash-linear-attention/SGLang before FP32 multiply. @@ -1321,6 +1322,32 @@ def chunk_kda( if initial_state is None or initial_state_indices is None: raise ValueError("KDA prefill requires an indexed initial-state pool") + num_tokens = q.shape[1] + if g.shape[1] < num_tokens or beta.shape[1] < num_tokens: + raise ValueError("g and beta must cover every q token") + g = g[:, :num_tokens] + beta = beta[:, :num_tokens] + if num_tokens == 1: + # Tracing constant-folds size-one dimensions, but the resulting kernel + # can share a cache entry with longer inputs. Keep T=1 on Triton so a + # short first request cannot specialize later Helion calls incorrectly. + return triton_chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + initial_state_indices=initial_state_indices, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + cu_seqlens=cu_seqlens, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + output_intermediate_states=output_intermediate_states, + ) + q = q.contiguous() k = k.contiguous() if use_qk_l2norm_in_kernel: diff --git a/test/registered/kernels/ops/attention/test_kda_helion.py b/test/registered/kernels/ops/attention/test_kda_helion.py index b2894cb61..36a8f34b3 100644 --- a/test/registered/kernels/ops/attention/test_kda_helion.py +++ b/test/registered/kernels/ops/attention/test_kda_helion.py @@ -28,6 +28,7 @@ else: ) from sglang.kernels.ops.attention.helion.kda_prefill import ( _intra_matrices_wide, + _l2norm_qk, ) from sglang.kernels.ops.attention.helion.kda_prefill import ( chunk_kda as helion_chunk_kda, @@ -851,6 +852,122 @@ def test_fixed_partial_prefill_and_state_pool_contract() -> None: assert torch.equal(helion_state[untouched], state[untouched]) +@pytest.mark.parametrize("is_varlen", [False, True], ids=["fixed", "varlen"]) +def test_single_token_prefill_does_not_poison_later_shapes( + is_varlen: bool, +) -> None: + """Keep a size-one first trace from specializing later prefill calls.""" + torch.manual_seed(991) + heads, key_dim, value_dim = 2, 32, 32 + a_log = torch.full([heads], -2.0, device="cuda") + dt_bias = torch.zeros(heads * key_dim, device="cuda") + indices = torch.zeros(1, device="cuda", dtype=torch.int32) + + _l2norm_qk.reset() + try: + for tokens in (1, 3): + q = torch.randn( + 1, tokens, heads, key_dim, device="cuda", dtype=torch.bfloat16 + ) + k = torch.randn_like(q) + v = torch.randn( + 1, tokens, heads, value_dim, device="cuda", dtype=torch.bfloat16 + ) + gate = torch.randn_like(q) * 0.2 + beta = torch.rand(1, tokens, heads, device="cuda") + state = ( + torch.randn( + 1, heads, value_dim, key_dim, device="cuda", dtype=torch.float32 + ) + * 0.01 + ) + cu_seqlens = ( + torch.tensor([0, tokens], device="cuda", dtype=torch.int32) + if is_varlen + else None + ) + + _compare_prefill( + q, + k, + v, + gate, + beta, + state, + indices, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + A_log=a_log, + dt_bias=dt_bias, + ) + finally: + _l2norm_qk.reset() + + +@pytest.mark.parametrize("is_varlen", [False, True], ids=["fixed", "varlen"]) +def test_prefill_ignores_padded_gate_rows(is_varlen: bool) -> None: + torch.manual_seed(997) + tokens, padded_tokens, heads, key_dim, value_dim = 51, 64, 2, 32, 32 + q = torch.randn(1, tokens, heads, key_dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, tokens, heads, value_dim, device="cuda", dtype=torch.bfloat16) + gate = -torch.rand( + 1, padded_tokens, heads, key_dim, device="cuda", dtype=torch.float32 + ) + beta = torch.rand(1, padded_tokens, heads, device="cuda") + gate[:, tokens:] = 1e4 + beta[:, tokens:] = 1e4 + cu_seqlens = ( + torch.tensor([0, 17, tokens], device="cuda", dtype=torch.int32) + if is_varlen + else None + ) + indices = ( + torch.tensor([0, 1], device="cuda", dtype=torch.int32) + if is_varlen + else torch.zeros(1, device="cuda", dtype=torch.int32) + ) + initial_state = torch.randn( + indices.numel(), + heads, + value_dim, + key_dim, + device="cuda", + dtype=torch.float32, + ) + + def run( + gate_input: torch.Tensor, beta_input: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + state = initial_state.clone() + output, chunks = helion_chunk_kda( + q, + k, + v.clone(), + gate_input, + beta_input, + initial_state=state, + initial_state_indices=indices, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + output_intermediate_states=True, + ) + return output, chunks, state + + trimmed = run(gate[:, :tokens], beta[:, :tokens]) + padded = run(gate, beta) + for padded_value, trimmed_value in zip(padded, trimmed): + assert torch.equal(padded_value, trimmed_value) + + short_inputs = ( + (gate[:, : tokens - 1], beta[:, :tokens]), + (gate[:, :tokens], beta[:, : tokens - 1]), + ) + for short_gate, short_beta in short_inputs: + with pytest.raises(ValueError, match="g and beta must cover every q token"): + run(short_gate, short_beta) + + def test_prefill_uses_stable_subchunk_gates() -> None: torch.manual_seed(1117) tokens, heads, key_dim, value_dim = 64, 1, 32, 32