[Feature] Intra-Card Context Parallelism to improve B=1 long-sequence performance

[Feature] Intra-Card Context Parallelism to improve B=1 long-sequence performance

Problem

FlashKDA's kernel2 (recurrence) launches with grid (N, H) — one block per head per sequence. At B=1 (e.g., single long-sequence prefill), SM utilization is critically low:

This causes FlashKDA to be slower than Triton baselines (fla_chunk_kda) at B=1:

Proposed Solution

Adapt FlashQLA's Intra-Card Context Parallelism strategy — split long sequences into sub-segments and process them in parallel:

1. Sequence splitting: Automatically compute optimal sub-segment length based on SM count and head count
2. Two-pass forward:
   • Pass 1: Run all sub-segments with h0=0 in parallel, capture each sub-segment's final_state
   • Pass 2: Chain final_states as initial_states for subsequent segments, re-run in parallel to produce correct output
3. Safety guarantee: Analytically estimate gate decay via A_log; only enable CP when initial state contribution is negligible

Advantages

• Pure Python-level implementation — no CUDA kernel modifications required
• Leverages existing cu_seqlens + initial_state interface
• Automatically decides whether to engage CP; zero overhead for batched scenarios

Start by tracing the kernel2 recurrence and the existing cu_seqlens and initial_state interfaces, then review FlashQLA's Intra-Card Context Parallelism strategy referenced in the proposal. Done means B=1 long-sequence cases can use safe sub-segment parallelism, preserve correct chained states, and avoid overhead for batched scenarios.