[glm] [pcp] fuse norm rope with pcp cache gather

[GirasoleY]

Summary

Fuse PCP Q/K preparation, cache quantization, and NVLS publication into one dispatch kernel plus one combine kernel for prefill, mixed, and decode batches.

• Old: bf16 norm/RoPE: output kv_c, k_pe, indexer_k → 3 all-gathers → quantize + cache-insertion
• New: fused dispatch + NVLS multicast → local cache combine/scatter

Dispatch

It launches one 256-thread CTA per local token and performs:

• Q RMSNorm, keeping Q local because each rank already holds the query rows it computes.
• KV RMSNorm.
• RoPE on k_pe.
• For sparse layers, indexer-K LayerNorm and RoPE.
• Packing directly into a cache-ready representation and NVLS multicast to every PCP rank.

For static E4M3, the fixed transport row is:

| kv_c E4M3: 512 B | k_pe E4M3: 64 B | indexer K E4M3: 128 B | index scale FP32: 4 B | alignment: 12 B | = 720 B/token

For fp8_ds_mla, the fixed transport row is:

| kv_c E4M3: 512 B | KV tile scales FP32: 16 B | k_pe BF16: 128 B | indexer K E4M3: 128 B | index scale FP32: 4 B | alignment: 12 B | = 800 B/token

Full workspace

workspace[double_buffer_index][src rank][max_num_tokens][payload]

Combine

The second kernel reads the received rank-major rows and scatters them into each rank's local paged caches:

• Prefill: tokens from every PCP rank are inserted.
• Decode: tokens are replicated, so they are inserted once from rank 0.

Scope and limitations

• Requires CUDA and NVLS symmetric-memory multicast.
• Specialized for the GLM-5.2 / DeepSeek-V3.2 MLA dimensions above.
• Each rank retains a persistent double-buffered staging allocation of approximately 2 * num_ubatches * PCP_world_size * max_local_tokens * packed_row_bytes.

Performance

GLM-5.2 PCP8 on one 8×B300 node, TP1/PCP8/EP8, BF16 activations, static E4M3 KV cache, a 32K batched-token limit, and a 32K long-prefill threshold.

Default portable PCP pipeline (historical stage profile):

Stage	Time
Fused Q/K norm + K RoPE	1.441 ms
Indexer-K all-gather	2.064 ms
Indexer-K quantization/insertion	0.423 ms
MLA kv_c all-gather	11.044 ms
MLA k_pe all-gather	2.349 ms
MLA quantization/insertion	6.931 ms
Total	24.252 ms

Final fused NVLS pipeline (latest two-run mean):

Stage	Time
Fused norm/RoPE + FP8 quantization + NVLS multicast	8.228 ms
Cache combine/scatter	0.898 ms
Total	9.127 ms

Controlled current-stack full-forward A/B from the Kineto execute_context_1(32768)_generation_0(0) GPU annotation:

Metric	Default portable PCP	Final fused NVLS	Change
Mean across PCP ranks	357.781 ms	343.875 ms	−3.89%
Slowest PCP rank	357.933 ms	344.109 ms	−3.86%

Each value is the arithmetic mean of two fresh-process runs. The portable baseline and pre-optimization fused path were measured in one current-stack ABBA; the final fused kernel was then measured in a second ABBA against that fused control. All runs used the same B300 node, model snapshot, source tree, eager-mode configuration, exact 32K input, warmup, and profiler procedure. The second ABBA measured the slowest-rank fused forward at 345.069 ms before these cache-overhead changes and 344.109 ms after them (−0.960 ms, −0.28%). "Default portable PCP" means the Triton fused norm/RoPE + PCP gather + cache-insert path.

Relationship to existing work

• [nv] add pcp support in dsv3.2 #52046 supplies the DeepSeek-V3.2 PCP gather-and-insert baseline now on main.
• [refactor] consolidate cp attn ops #52839 moved the PCP attention helpers into vllm/v1/attention/ops/pcp.py; this branch is ported onto that layout.
• [Kernel][PCP] Publish replicated KV updates through PyTorch symmetric memory #52863 writes final rows directly into symmetric persistent cache replicas through P2P peer pointers.
• [Kernel][Core][PCP][KV Update] Replace PCP update collectives with direct peer stores #49517 has the same replicated-KV-update goal using CUDA VMM peer mappings and direct peer stores.

This PR uses a different design from #52863/#49517: NVLS multicast into transient symmetric staging followed by a receiver-local paged-cache combine. It keeps cache tensors, allocation, and ownership ordinary and receiver-local, supports PCP2/4/8, and integrates with the common DeepSeek-V3.2 PCP path.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @GirasoleY.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork