Merge branch 'main' of https://github.com/dylanyunlon/project_6

docs/CCCL_REDUCE_ARCHITECTURE_NOTES.md

@@ -0,0 +1,159 @@
+# CCCL Reduce Architecture Notes
+
+> Source: `dispatch_reduce.cuh`, `kernel_reduce.cuh`, `agent_reduce.cuh`, `tuning_reduce.cuh`, `util_arch.cuh`
+> Read: 2026-08-04 by Claude from CCCL upstream in project_6/cccl_upstream/
+
+## Key Architecture
+
+### Two-pass dispatch (dispatch_reduce.cuh)
+
+```
+num_items <= single_tile.threads * single_tile.items
+  → SingleTile: one CTA, one kernel launch
+  → DeviceReduceSingleTileKernel(d_in, d_out, num_items, ...)
+
+num_items > single_tile threshold
+  → Pass 1: DeviceReduceKernel — N CTAs each reduce their share → d_block_reductions[N]
+  → Pass 2: DeviceReduceSingleTileKernel — 1 CTA reduces d_block_reductions[N] → d_out
+```
+
+Grid size for Pass 1: `max_blocks = sm_occupancy * sm_count * subscription_factor(5)`
+For BI-V100: `2 * 16 * 5 = 160 blocks` max.
+Each block processes `ceil(num_items / 160)` elements.
+
+### Tile consumption (agent_reduce.cuh)
+
+**Critical: tile data is in registers, NOT SMEM.**
+
+```cpp
+AccumT items[ITEMS_PER_THREAD];  // <-- register array, per-thread
+// ... load from global memory ...
+thread_aggregate = ThreadReduce(items, reduction_op);  // per-thread reduction
+
+// Only SMEM used:
+BlockReduce(temp_storage.reduce).Reduce(thread_aggregate, reduction_op);
+```
+
+`TempStorage` = `BlockReduce::TempStorage` ≈ threads * sizeof(AccumT) bytes.
+NOT threads * items * sizeof(AccumT).
+
+### Vectorized loads
+
+```cpp
+ATTEMPT_VECTORIZATION = (vec_size > 1) && (ITEMS_PER_THREAD % vec_size == 0)
+    && is_pointer<InputIteratorT>
+    && (is_primitive<InputT> || is_trivially_relocatable<InputT>)
+    && sizeof(InputT) <= 8;
+```
+
+For fp32 scores: vec_size=2 → loads 8 bytes (2 floats) per instruction.
+For fp16 KV cache: vec_size=4 → loads 8 bytes (4 halfs) per instruction.
+
+### scale_mem_bound vs scale_reg_bound (util_arch.cuh)
+
+Two scaling functions with different constraints:
+
+**scale_mem_bound** (memory-bound algorithms: reduce, transform):
+- items = clamp(nominal * 4 / type_size, 1, nominal * 2)  ← allows 2x expansion
+- threads = min(nominal, round_up(48KB / (type_size * items), 32))
+
+**scale_reg_bound** (register-bound algorithms: scan with complex state):
+- items = max(1, nominal * 4 / max(4, type_size))  ← no expansion past nominal
+- threads = min(nominal, ceil_div(48KB / (type_size * items), 32) * 32)
+
+Key difference: scale_reg_bound uses `max(4, type_size)` preventing items from exceeding nominal for small types, and uses `ceil_div` instead of `round_up` for thread count. Both use 48KB as the cap, but this limits REGISTER PRESSURE (spill to local memory), not actual SMEM usage.
+
+## Impact on muh tuning
+
+### Our SMEM model was wrong for reduce
+
+`test_smem_safety.py` and `check_smem()` in `muh_kernel_map.py` compute
+`tile_bytes = threads * items * type_size` and check against 49152.
+
+This is the scale_mem_bound cap, NOT the actual SMEM usage. The actual SMEM
+for reduce is approximately `threads * max(sizeof(AccumT), 4)` bytes — about
+2-8 KB, not 32-49 KB.
+
+CCCL's SM100 float64 tuning uses `threads=640, items=16` → scale_mem_bound
+"tile" = 640*16*8 = 81920 > 49152. But this doesn't overflow SMEM — it only
+means scale_mem_bound will cap threads down. The actual kernel SMEM usage
+with threads=640 is only ~5120 bytes.
+
+### Our float64/int64 tuning may be too conservative
+
+We use threads=384 items=16 for float64, capped by scale_mem_bound. CCCL
+uses threads=640 items=16 on SM100. The question is whether BI-V100's register
+file (255 regs/thread) can hold 16 float64 items without spilling.
+
+16 * 8 = 128 bytes = 32 registers per thread for tile data alone.
+With overhead (thread_aggregate, loop variables, etc.), ~40 registers/thread.
+255 max registers → no spill risk. threads=640 may be safe on BI-V100.
+
+**TODO**: Benchmark threads=640 items=16 for float64 on BI-V100.
+
+### paged_attn.py forces V1
+
+Line 99: `use_v1 = True` overrides V1/V2 heuristic. V2 is completely disabled.
+For 100K token sequences, V1 makes one CTA iterate over all KV blocks — bad
+for latency. V2 would partition the work and reduce across partitions, which
+is exactly CCCL's two-pass pattern.
+
+**TODO**: Re-enable V2 for max_seq_len > 8192. Use muh's partition_size tuning.
+
+### _PARTITION_SIZE = 512 is hardcoded
+
+Not controlled by muh. Should be tunable: larger partition = fewer blocks =
+less overhead but more work per block. Optimal value depends on SM count.
+For 16 SMs: partition_size=1024 may be better (fewer partitions to reduce).
+
+---
+
+## CCCL Scan Architecture (dispatch_scan.cuh)
+
+> Added: 2026-08-04
+
+### Two algorithm paths
+
+**Lookback** (all GPUs including BI-V100):
+- Each CTA processes one tile, uses `ScanTileState` in global memory for inter-CTA communication
+- Lookback delay policy controls how aggressively CTAs poll predecessors
+- SMEM: static only (`__shared__`), passed as `0` dynamic SMEM
+- BI-V100 optimal: `no_delay` (dcid=0) because 16 SMs → ~32 CTAs → tile_status fits in 6MB L2
+
+**Lookahead** (SM100+ only, PTX ISA >= 860):
+- Pipeline-based with `__pipeline_memcpy_async` and bulk copy
+- Uses dynamic SMEM with auto-selected `num_stages`
+- **Not available on BI-V100** — requires NVIDIA PTX ISA 860+ instructions
+- All lookahead structs in our tuning_scan.cuh can remain empty shells
+
+### ScanTileState allocation
+
+Scan requires `d_temp_storage` for tile status descriptors:
+```
+tile_size = threads * items
+num_tiles = ceil(num_items / tile_size)
+temp_bytes = tile_state.AllocationSize(num_tiles)
+```
+
+For BI-V100 with 100K tokens and tile_size=384*22=8448:
+num_tiles = ceil(100000/8448) = 12 tiles → negligible temp storage.
+
+### Grid size for scan
+
+Lookback scan launches `num_tiles` blocks (one per tile), NOT `sm_count * subscription_factor`.
+This is different from reduce, which uses `GridEvenShare`.
+For scan, every CTA processes exactly one tile and communicates with neighbors.
+
+With 12 tiles on 16 SMs: all tiles fit in one wave, zero lookback contention.
+This is why `no_delay` works on BI-V100 — the entire scan completes in a single wave.
+
+### Lookahead num_stages optimization (SM100 only)
+
+CCCL dynamically selects pipeline depth:
+```cpp
+max_stages = ceil(num_items / (sm_count * tile_size)) + 1
+while (smem_for_stages(num_stages+1) <= max_dynamic_smem) num_stages++
+```
+
+For BI-V100 this is irrelevant (no pipeline support), but the formula shows
+NVIDIA's strategy: match pipeline depth to problem size / SM count ratio.

muh/gen_patch.py

@@ -268,29 +268,55 @@ def generate_patches(header_dir):
                 summary.append(f"SKIP {algo}: no bi100_* structs and no inline values found")
                 continue
 
-        # Select the most relevant struct for vllm's primary data path.
+        # Select the struct that matches each vllm kernel's data type.
-        # vllm's paged_attention score accumulator is always float32 (4 bytes),
+        #
-        # so we prefer bi100_*float32* or bi100_*accum4* structs.
+        # CCCL's policy_selector dispatches by (accum_size, type_t, offset_size).
-        # Fallback priority: float32 > accum2 (fp16 KV) > first non-default > first.
+        # gen_patch must do the same: when injecting into paged_attention
-        primary = None
+        # (float32 scores), use bi100_plus_float32_o4, not bi100_plus_accum1_o4.
-        preference_order = ['float32', 'accum4', 'accum2', 'int32']
+        #
-        for pref in preference_order:
+        # The VLLM_KERNEL_MAP in muh_kernel_map.py defines each kernel's
-            for name, fields in structs:
+        # data_types. This mapping encodes the primary data type per algorithm:
-                if pref in name and 'det' not in name and 'default' not in name:
+        ALGO_PRIMARY_TYPE = {
-                    primary = (name, fields)
+            'reduce':       ('float32', 4),   # paged_attention scores
-                    break
+            'scan':         ('float32', 4),   # softmax denominator
-            if primary:
+            'topk':         ('float32', 4),   # logits
-                break
+            'transform':    ('float16', 2),   # activations (SiLU, RMSNorm input)
-        if primary is None:
+            'batch_memcpy': ('float16', 2),   # KV cache blocks
-            for name, fields in structs:
+            'for':          ('float16', 2),   # RoPE
-                if 'default' not in name and 'det' not in name:
+        }
-                    primary = (name, fields)
-                    break
-        if primary is None:
-            primary = structs[0]
 
-        pname, pfields = primary
+        target_type, target_size = ALGO_PRIMARY_TYPE.get(algo, ('float32', 4))
-        summary.append(f"READ {algo}: {pname} → {pfields}")
+
+        # Score each struct by match quality
+        def struct_score(name, fields):
+            score = 0
+            name_lower = name.lower()
+            # Exact type name match (best)
+            if target_type.replace('float', 'f') in name_lower or target_type in name_lower:
+                score += 100
+            # Accum/type size match in name (e.g. "_4B_", "_accum4_", "float32")
+            size_tags = [f'_{target_size}B', f'_accum{target_size}', f'float{target_size*8}']
+            for tag in size_tags:
+                if tag.lower() in name_lower:
+                    score += 50
+            # Offset size 4 preferred (most common in vllm)
+            if '_o4' in name_lower:
+                score += 10
+            # Penalize 'default' and 'det' (deterministic) structs
+            if 'default' in name_lower:
+                score -= 200
+            if 'det' in name_lower:
+                score -= 50
+            # Penalize 1-byte type structs for float32 targets
+            if target_size >= 4 and ('_1B' in name or 'accum1' in name_lower):
+                score -= 100
+            return score
+
+        scored = [(struct_score(n, f), n, f) for n, f in structs]
+        scored.sort(key=lambda x: -x[0])
+        _, pname, pfields = scored[0]
+
+        summary.append(f"READ {algo}: {pname} → {pfields} (target: {target_type})")
 
         for field_name, value in pfields.items():
             key = (algo, field_name)