Add prebuilt metadata support and optimize conv1d operations

[AndyLi429]

Summary

This PR removes the scalar device-to-host (D2H) synchronizations that the NPU
kernels issue on the Qwen3.5 GDN (gated delta net) prefill path —
query_start_loc[-1], seqlens.max(), has_initial_state.any(), and the
equivalents in the FLA chunk kernels — by letting the caller pass the host-side
metadata that the scheduler already holds on the CPU. Eliminating these
per-layer scalar reads keeps the NPU stream from stalling during prefill.

It is a dependency of the SGLang change
sgl-project/sglang#24597: this kernel PR must be merged
first, then the SGLang side.

The change has two parts plus a test.

1. Prebuilt metadata and CPU sequence lengths for the FLA chunk path

chunk_gated_delta_rule_npu (and every kernel it calls — chunk_local_cumsum,
chunk_scaled_dot_kkt_fwd, solve_tril, recompute_w_u_fwd,
chunk_gated_delta_rule_fwd_h, chunk_fwd_o) now accept two optional inputs:

• cu_seqlens_cpu: a CPU copy of the cumulative sequence lengths, used to build
  chunk_indices / chunk_offsets without reading scalars off a device tensor.
• prebuilt_meta: pre-computed chunk_indices / chunk_offsets variants for the
  different block sizes (chunk64, large_block). When supplied, the kernels
  skip recomputing them on every call.

Each kernel follows the same fallback order: use the prebuilt value if given,
otherwise compute from cu_seqlens_cpu, otherwise fall back to cu_seqlens
(legacy D2H path). A guard in chunk_gated_delta_rule_npu clears
cu_seqlens_cpu / prebuilt_meta when cu_seqlens is None to prevent silent
misuse in non-varlen mode, and the varlen output is trimmed to the actual
sequence length via the host cu_seqlens_cpu[-1].

Also fixes a dtype bug in chunk_local_cumsum_scalar_kernel: the result was
stored as s.dtype (input dtype) instead of o.dtype (output dtype).

2. Host-side input contract for causal_conv1d_fn_npu

The varlen prefill wrapper previously read three scalars off device tensors on
every call — query_start_loc[-1], seqlens.max(), has_initial_state.any() —
each forcing a D2H sync. It now accepts host-side equivalents and only falls back
to the device reads when a (legacy) caller omits them:

• seq_lens_cpu: when query_start_loc is None, the wrapper builds the
  cumulative query_start_list from this host list (asserting it is a host list /
  CPU tensor, never a device tensor), then derives cu_seq_len, max_query_len
  and batch_size on the CPU.
• cu_seq_len / max_query_len: host scalars that replace query_start_loc[-1]
  and seqlens.max() inside prepare_data.
• has_initial_state_any: host boolean that replaces has_initial_state.any()
  when deciding whether to gather the initial states.

Null-safety / ergonomics:

• prepare_data guards on conv_states is not None and cache_indices is not None
  before indexing, fixing a crash when either is None.
• cache_indices defaults to arange(batch_size) when not provided.
• causal_conv1d_fn_npu and causal_conv1d_update_v2 keep a **kwargs
  catch-all so an older SGLang caller — or a newer caller against an older
  wrapper — does not raise TypeError. This is the cross-version back-compat
  anchor between sgl-kernel-npu and SGLang.

3. Tests

test_causal_conv1d_fn_prefill_uses_cpu_seq_lens (in test_mamba_conv.py)
verifies that the new seq_lens_cpu path (query_start_loc=None, lengths given
as a host list) produces the same output and the same updated conv_states
as processing each sequence individually.

[gemini-code-assist]

Code Review

This pull request introduces support for prebuilt metadata and CPU-side sequence lengths (cu_seqlens_cpu) across several NPU kernels in the fla module to optimize performance. The changes involve updating kernel wrappers to accept pre-calculated indices and offsets and ensuring fallback logic utilizes CPU sequence lengths when provided. Feedback was provided to address several critical issues where newly generated indices were not explicitly moved to the NPU device, which would lead to device mismatch errors, and to fix a potential TypeError in the main entry point when sequence lengths are not provided.

[AndyLi429]

/gemini review

[gemini-code-assist]

Code Review

This pull request optimizes several NPU kernels in the fla and mamba modules by passing CPU-side sequence length information and pre-built metadata to avoid GPU-to-CPU synchronization. Key changes include updating chunk_gated_delta_rule, chunk_local_cumsum, and causal_conv1d_fn_npu to utilize these CPU tensors for indexing and offsets. The PR also fixes a bug in the cumsum kernel where the store operation used an incorrect element type and adds output slicing to ensure correct tensor shapes. I have no feedback to provide.