[ENGINE] apply CCCL CacheAsyncConfiguration pattern to activation/layernorm
Source: cccl_upstream/cub/cub/device/dispatch/dispatch_transform.cuh
(CacheAsyncConfiguration + spread_out_items_per_thread)
CCCL dispatch_transform.cuh insight: element-wise transforms have
deterministic output shapes. Cache output tensors to avoid cudaMalloc.
Quote from CCCL: 'This computation MUST NOT depend on runtime state
... since the result will be cached.'
Applied to:
1. GeluAndMul.forward_cuda — output tensor cached during decode
2. RMSNorm.forward_cuda — output tensor cached during decode
(64 layers × 2 norms/layer = 128 cudaMalloc eliminated per step)
SiluAndMul already had this pattern from previous commit.
BI-V100 has no async memory allocator — synchronous cudaMalloc blocks
the entire SM pipeline. Eliminating 128+ allocations per decode step
directly improves Output TPS (83% competition weight).
This commit is contained in:
@@ -97,7 +97,22 @@ class GeluAndMul(CustomOp):
|
||||
|
||||
d = x.shape[-1] // 2
|
||||
output_shape = (x.shape[:-1] + (d, ))
|
||||
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# CCCL dispatch_transform.cuh CacheAsyncConfiguration pattern:
|
||||
# Output tensor shape is deterministic from input shape.
|
||||
# During decode, shapes are stable → cache to avoid cudaMalloc.
|
||||
# CCCL: "This computation MUST NOT depend on runtime state ...
|
||||
# since the result will be cached."
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
_cache_key = (output_shape, x.dtype, x.device)
|
||||
_cached = getattr(self, '_out_cache', {}).get(_cache_key)
|
||||
if _cached is not None and _cached.shape == output_shape:
|
||||
out = _cached
|
||||
else:
|
||||
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
||||
if not hasattr(self, '_out_cache'):
|
||||
self._out_cache = {}
|
||||
self._out_cache[_cache_key] = out
|
||||
if self.approximate == "none":
|
||||
ops.gelu_and_mul(out, x)
|
||||
elif self.approximate == "tanh":
|
||||
|
||||
@@ -85,7 +85,28 @@ class RMSNorm(CustomOp):
|
||||
residual_alpha,
|
||||
)
|
||||
return x, residual
|
||||
out = torch.empty_like(x)
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
# CCCL dispatch_transform.cuh CacheAsyncConfiguration pattern:
|
||||
# Element-wise transforms have deterministic output shapes.
|
||||
# During decode, input shape is stable (num_seqs × hidden_dim).
|
||||
# Cache the output tensor to avoid cudaMalloc on every step.
|
||||
#
|
||||
# CCCL: "This computation MUST NOT depend on runtime state ...
|
||||
# since the result will be cached."
|
||||
#
|
||||
# RMSNorm is called 64× per forward pass (Qwen3.6 has 64 layers).
|
||||
# Each call was doing torch.empty_like → cudaMalloc.
|
||||
# With caching: 64 cudaMalloc calls → 0 per decode step.
|
||||
# ═══════════════════════════════════════════════════════════════
|
||||
_cache_key = (x.shape, x.dtype, x.device)
|
||||
_cached = getattr(self, '_out_cache', {}).get(_cache_key)
|
||||
if _cached is not None and _cached.shape == x.shape:
|
||||
out = _cached
|
||||
else:
|
||||
out = torch.empty_like(x)
|
||||
if not hasattr(self, '_out_cache'):
|
||||
self._out_cache = {}
|
||||
self._out_cache[_cache_key] = out
|
||||
ops.rms_norm(
|
||||
out,
|
||||
x,
|
||||
|
||||
Reference in New Issue
Block a user