[muse_spark][perf] Split the packed qkv inside the q/k region

[felipemello1]

TLDR: Muse Spark's muse_spark_attention_qk region now takes the packed wqkv output and splits it itself. The q/k copies fuse into the per-head norms, and the backward writes the packed grad without a cat. One 4-layer unit: decode -0.9% to -2.4%, prefill -0.4% to -1.1%, train 16k -0.14% by kernel time (wall rows within noise). Outputs and grads are bitwise identical.

Problem

1. Packed projection. QKVLinear.wqkv returns q, k and v as one tensor, viewed as [T, 16 kv heads, 10, 64]: per kv head, 8 q heads + 1 k + 1 v.
2. Split in eager. QKVLinear.forward splits it and copies q, k and v out (q's reshape can't be a view; vLLM needs contiguous k and v).
3. Read again. muse_spark_attention_qk then reads the copied q and k for the norms and RoPE.
4. Backward. Autograd cats the three grads back into one packed grad, in eager.
5. Cost. Per 4-layer unit: train 16k 0.79 of 333 ms (copies 0.41 + cat 0.38); prefill 16k 0.41 of 102 ms; decode T=1 12 copy kernels, 28 us of 1.84 ms.

Solution

1. Split inside the region. MuseSparkAttention calls qkv_linear.wqkv(x) and passes the packed tensor to the region, now _prepare_query_key_value, which returns q, k and v.
2. QKVLinear.split. The split moves out of QKVLinear.forward into its own method, with the same local_map output types; the region calls it. wqkv now runs in global SPMD mode, as Muse's output_gate already does (SPMD typecheck fwd+bwd passes before and after for Muse Spark, Muse Glimmer, Llama 3 and GPT-OSS at TP=2 with and without SP, and Muse Spark and Llama 3 at CP=2).
3. What fuses. The q and k copies fuse into the norm kernels. The cat fuses into the norm-backward kernels, which write the packed grad. v is still copied, now by Inductor.
4. Unchanged: QKVLinear.forward runs the same ops for Llama 3, Qwen3, GPT-OSS and Muse Glimmer; q, k and v keep their contiguous layouts; the region's options and graph count.

before:  wqkv -> split + 3 copies (eager) -> region(q, k): norm -> RoPE     bwd: region bwd -> cat (eager)
after:   wqkv -> region(qkv): split -> norm -> RoPE, copy v                 bwd: region bwd writes the packed grad

Results

Numerics: bitwise identical to before (GB300). q, k and v of a RoPE layer and a NoPE layer, train and no_grad,
          the wqkv weight grads and the input grad. The change only moves copies.

Muse Spark 816B, layers 0-3 (3 sliding + RoPE, 1 global), native dims, TP=1, one EP32 rank's experts (stock),
no activation checkpointing, flex attention, CUDA graphs, 1x GB300, ms per unit, min of 2 interleaved processes
mode             T      before    after    delta
train fwd+bwd    2k     34.32     33.92    -1.2%
                 4k     71.32     70.96    -0.5%
                 16k   333.39    331.92    -0.4%
prefill no_grad  2k     11.45     11.32    -1.1%
                 16k   101.58    101.19    -0.4%
decode no_grad   1      1.839     1.794    -2.4%
                 8      1.865     1.849    -0.9%
                 64     2.256     2.231    -1.1%
                 256    2.402     2.362    -1.7%
train 16k kernel time: split, copies, cat and region kernels 1.69 -> 1.21 ms per unit (-0.48 ms, 0.14%); the larger
wall delta is noise.

Glue alone (split, copies, q/k norms, RoPE; GEMMs excluded), layers 0 + 1, kernel time, us
train 16k     1074 -> 824
prefill 16k    504 -> 340
decode 1      23.9 -> 12.1   (12 -> 8 kernels)

graphs: muse_spark_attention_qk 8 after train 2k/4k/16k, prefill 2k/16k, decode 1/8/64/256 (limit 16), as before
other QKVLinear users (Llama 3 8B, Qwen3 8B, GPT-OSS 20b, Muse Glimmer 30B), qkv_linear fwd+bwd 4k/16k and
decode 1: same kernels; min of 3 processes within -0.2%..+0.9% (16k varies 1427-1494 us within each tree)

You may ask: why not Llama 3, Qwen3 and GPT-OSS too? Their RoPE region is the RoPE module's own forward, which takes q and k. Fusing the split needs a new q/k region in GQAttention (split -> norm -> RoPE): a follow-up. The copies and cat cost Qwen3 8B 127 us per layer at 16k.
You may ask: why is training's q norm kernel slower (114 vs 76 us per layer at 16k)? It also writes the split q, which the norm backward needs; before, the eager copy provided that tensor.
NOTE: reaches every Muse Spark recipe (the region is on by default): the training recipes and the RL DAPO-math recipe (Dist-MoE experts, so the microseconds carry over, the percentages don't).
NOTE: not tested: loss curve, multi-GPU training runs.

Test plan
pytest tests/unit_tests/gpu/models/muse_spark/test_attention_compile.py (the q/k region now takes packed qkv; checks q, k, v and the packed grad against eager)
pytest tests/unit_tests/cpu/models/common/test_qkv_linear.py
pytest tests/unit_tests/gpu/models/common/test_async_linear.py (2 GPUs; the async wqkv now runs under global type checking)
SPMD typecheck on a 2-rank CPU mesh, fwd+bwd, before and after (one-off checks): Muse Spark (RoPE and NoPE layers), Muse Glimmer, Llama 3 and GPT-OSS (bias) at TP=2 with and without SP; Muse Spark and Llama 3 at CP=2.