refactor(kv): index_copy_ for KV-cache writes (enabler for multi-user batching + CUDA Graphs)

[prashant182]

Summary

Replaces kv_cache[\"k\"][:, current_start:current_end] = roped_key with kv_cache[\"k\"].index_copy_(1, kv_write_index, roped_key) in the local_attn_size == -1 fast-path. kv_write_index is a [seq_lens]-shape tensor built once per chunk in generate() and threaded as a kwarg. Mirrored in both non-SP (CausalWanSelfAttention.forward) and SP (sp_attn_forward_causal) paths.

Bit-identical at B=1 (MD5 ed2f82628308a3f8acd9b7935bb84401), works correctly at B=2 (verified in an isolated test).

What this is and isn't

This PR ships no perf win on its own at B=1 — it's a refactor that swaps one correct primitive for another. The reason it's worth landing now is that it opens two doors that the slice-indexing form keeps shut:

1. Multi-user batching (B>1). Slice-assign with a Python-int current_start works at B=1 but doesn't generalize to multi-user inference, where each batch element may be writing the same positions of its own KV slab. index_copy_ does generalize — the same kv_write_index is correct for all batch elements in homogeneous batching (different users on the same chunk timeline). Heterogeneous batching needs PagedAttention; explicitly out of scope here.
2. CUDA Graphs. Removes one of two Python-int slice sources that were keeping the DiT forward out of torch.compile(mode='reduce-overhead'). The cache READ slice (cache[\"k\"][:, :local_end_index]) is still Python-int; addressing that is a separate experiment.

Verification

• Isolated test mirrors CausalWanSelfAttention.forward's control flow with tiny dims (no FSDP/SP). Three cases:
  • B=1, 7-chunk sequence — outputs and final cache state bit-equal to slice path.
  • B=2, 7-chunk sequence — outputs and final cache state bit-equal (validates the multi-user precondition).
  • torch.compile probe — dynamic=True collapses unique graph count from 5 → 2 (read-slice work pending for a follow-up).
• End-to-end: bench MD5 ed2f82628308a3f8acd9b7935bb84401 matches the locked baseline; generate() 13414 ms is within run-to-run noise of the post-B3 baseline (13523 ms).

Non-regressive design

Every new kwarg (kv_write_index) defaults to None. With None, the slice path remains the eager fallback. External callers of WanModelFast.forward, sp_dit_forward_causal, or the attention forwards keep working unchanged.

Stack

Stacked on #51 (B3). Reported diff shrinks once that lands.

[prashant182]

Nudge for review when you have a moment please cc @Robbyant @JingyeChen @qiuyu96