moe_align_block_size() allocates 3 tensors per call:
sorted_ids (int32, ~320 elements for decode)
expert_ids (int32, ~320 elements)
num_tokens_post_pad (int32, 1 element)
Called 64 times per decode step (once per MoE layer) = 192 CUDA mallocs.
During decode, these shapes are stable (same num_seqs × topk × num_experts).
Fix: cache in _moe_intermediate_cache (same dict as intermediate_cache1/2/3).
Reuse when shapes match. First call allocates, subsequent 63 calls reuse.
Combined with d3b1108 (intermediate cache): total savings = 189 + 192 = 381
CUDA mallocs eliminated per decode step.
At 395 TPS target: 381 × 395 = 150,495 fewer mallocs/second.
CCCL source read as input: tuning_transform.cuh (549 lines)
Key insight extracted: cc_to_min_bytes_in_flight maps hardware to prefetch
depth. BI-V100 = 64KB (B200 level). But more importantly, the policy_selector
architecture shows that the dispatch layer (Python) should minimize overhead
to let the kernel layer (C++/ixformer) run uninterrupted — which is exactly
what tensor pre-allocation achieves.
fused_experts() is called 64 times per decode step (once per MoE layer).
Each call allocated 3 intermediate tensors via torch.empty = 192 mallocs.
For decode (M=1, topk=8), all 64 calls use identical shapes.
Fix: module-level _moe_intermediate_cache dict that reuses tensors when
shapes match. First layer call allocates, subsequent 63 calls reuse.
Saves 189 CUDA mallocs per decode step = 74,655 mallocs/second at 395 TPS.
Design follows CCCL's dispatch_reduce.cuh pattern: pre-allocate temp_storage
once via alias_temporaries, reuse across kernel invocations.
No functional change — tensors are .empty() (uninitialized), overwritten
before use by ixformer kernels.
CCCL saxpy.cu demonstrates the principle: fused operations should minimize
wasted work. The saxpy_fast (single transform) vs saxpy_slow (two transforms)
comparison shows that eliminating unnecessary memory round-trips is the
primary optimization lever for element-wise ops.
Applied to MoE: during decode, M=8 (max-num-seqs) × topk=8 = 64 tokens.
Old heuristic: numel≤64 → BLOCK_SIZE_M=32 → 2 tiles of 32, no waste.
But for smaller batches (M=1,2,4 × topk=8 = 8,16,32 tokens):
BLOCK_SIZE_M=32 → tile padding: 24/16/0 rows wasted per tile
BLOCK_SIZE_M=16 → tile padding: 8/0/0 rows wasted per tile
New heuristic adds a finer-grained tier:
numel ≤ 16 → BLOCK_SIZE_M = 16 (zero waste for ≤2 seqs)
numel ≤ 64 → BLOCK_SIZE_M = 32 (was: same, no change)
numel ≤ 1024 → BLOCK_SIZE_M = 64 (was: same, no change)
else → BLOCK_SIZE_M = 256 (was: same, no change)
ixformer only reads BLOCK_SIZE_M from the config dict. The 16→32 threshold
matters for low-batch decode on BI-V100 where 16 SMs benefit from more
tiles with less padding over fewer tiles with more padding.
Source: cccl_upstream/thrust/examples/saxpy.cu (fusion + waste minimization)
