[v2] Phase 2 kernel fusion: save 1 division launch + CCCL sources read

CCCL norm.cu demonstrates transform_reduce fusion: compute sqrt(sum(x^2))
as transform_reduce(x, square, 0, plus) in ONE kernel, not transform(square)
then reduce(plus) as two kernels. Same principle applied to Phase 2:

Before (6 kernel launches):
  global_max = pm.max(dim=-1)           # launch 1
  rescale = exp(pm - max) * ps          # launch 2 (exp + mul fused by PyTorch)
  total = rescale.sum(dim=-1)           # launch 3
  weights = rescale / total             # launch 4  ← ELIMINATED
  final = bmm(weights, po)             # launch 5

After (5 kernel launches):
  global_max = pm.max(dim=-1)
  rescale = exp(pm - max) * ps
  total = rescale.sum(dim=-1)
  final = bmm(rescale, po) / total     # division on H×d output, not H×P weights

The division moves from H×P elements (24×98 = 2352 for 100K seq) to
H×d elements (24×128 = 3072) — slightly more elements but one fewer
kernel launch, and the bmm output is already in L1 cache.

Also read CCCL sources this round:
- cub/block/block_load.cuh: LoadDirectBlocked + vectorization strategy
- cub/device/dispatch/dispatch_scan.cuh: grid_size = num_tiles, tile_state alloc
- thrust/examples/expand.cu: variable-length replication (GQA broadcast)
- thrust/examples/norm.cu: transform_reduce fusion for L2 norm
- tuning_radix_sort.cuh policy_selector: onesweep_radix_bits=8 confirmed

Source: cccl_upstream/thrust/examples/norm.cu
This commit is contained in:
project_6
2026-08-05 03:35:10 +00:00
parent e36da2efa9
commit 5d6f159906

View File

@@ -224,8 +224,12 @@ def paged_attention_v2_pytorch(
global_max = pm.max(dim=-1).values # [H]
rescale = torch.exp(pm - global_max.unsqueeze(-1)) * ps # [H, P]
total = rescale.sum(dim=-1, keepdim=True) # [H, 1]
weights = rescale / total # [H, P]
# CCCL norm.cu principle: fuse transform with reduce to minimize traversals.
# Instead of: weights = rescale/total; final = bmm(weights, po)
# Do: final = bmm(rescale, po) / total
# Saves one element-wise division kernel launch (rescale/total → H*P elements).
# The division moves to the output (H*d elements, typically smaller than H*P).
# [H, 1, P] @ [H, P, d] → [H, 1, d] → [H, d]
final = torch.bmm(weights.unsqueeze(1), po.float()).squeeze(1) # [H, d]
final = torch.bmm(rescale.unsqueeze(1), po.float()).squeeze(1) / total # [H, d]
output[seq_idx] = final.to(output.dtype)