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:

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