fix(kernel) Fix Helion small-token prefill bug (#35197)

Co-authored-by: Ethan Che <eche@meta.com>

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} "

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: