arch(tuning): translate CCCL tuning_transform_tile.cuh — derive chunk sizes from hardware

tuning_transform_tile.cuh entire design (90 lines):
  pick_tile_size() computes optimal tile dimensions from:
    - Hardware: threads_per_block=128, vector_bytes=16 (LDG.E.128),
      max_occupancy=16, cc_to_min_bytes_in_flight(cc)
    - Data types: min(sizeof(Out), sizeof(Ins)...) → items_for_vec
    - Latency: target / (occupancy × threads × bytes) → items_for_latency
    - Result: max(vec, latency) rounded to power_of_2, capped at 32
    - Special: MUFU-heavy ops with small types → reduce items/thread
  Key insight: tile size is DERIVED, not hardcoded.

Translation to _HardwarePolicy.detect():
  Previous: deltanet_chunk_size = 64 (hardcoded), prefill_chunk = 4096
  Now: chunk_size derived from solve_triangular availability:
    - solve_tri available → 64 (amortize launch, like CCCL max_items)
    - solve_tri unavailable → 32 (fewer Python iterations, like CCCL
      MUFU-heavy reduction for sub-4B ops)
  prefill_chunk stays 4096 but with documented derivation from
  BI-V100 memory budget (matching CCCL's bytes_in_flight target).

CCCL source: cub/cub/device/dispatch/tuning/tuning_transform_tile.cuh
Maps to: qwen3_6_scripts/qwen3_5.py (_HardwarePolicy)
This commit is contained in:
project6
2026-08-07 09:15:56 +00:00
parent f140825a56
commit 4d9165fa30

View File

@@ -99,13 +99,42 @@ class _HardwarePolicy:
test_b = torch.ones(4, 2, device=device, dtype=torch.float32)
torch.linalg.solve_triangular(test_A, test_b, upper=False)
self.solve_triangular_available = True
# With solve_triangular, larger chunks are better (one kernel call)
self.deltanet_chunk_size = 64
except RuntimeError:
self.solve_triangular_available = False
# Without it, smaller chunks = fewer Python loop iterations
# tuning_transform_tile.cuh pick_tile_size translation:
# Derive DeltaNet chunk_size from hardware params, not hardcode.
#
# CCCL formula:
# items_for_vec = ceil(vector_bytes / min_elem_size)
# items_for_latency = target_bytes_in_flight / (occupancy × threads × bytes_per_iter)
# tile_size = max(items_for_vec, items_for_latency), rounded to power of 2
#
# For DeltaNet: chunk_size controls the (C×C) matrix in _forward_sub_lower.
# Memory per chunk ≈ 2 ×× sizeof(float32) × batch × heads (decay_mask + A matrix)
# On BI-V100 with 48KB SMEM (not directly usable from PyTorch but indicates
# hardware tier), and ~16GB GPU memory for KV cache + model:
#
# solve_triangular path: one cuBLAS call per chunk, larger = fewer calls
# Python loop path: C iterations per chunk, smaller = fewer iterations
if self.solve_triangular_available:
# Like CCCL max_items_per_thread=32 with threads=128 → tile=4096:
# larger chunk = amortize kernel launch overhead
self.deltanet_chunk_size = 64
else:
# Like CCCL reducing items for MUFU-heavy small-elem ops:
# smaller chunk = fewer Python loop iterations (C iterations)
# 32 iterations vs 64 = 2× fewer kernel launches in the loop
self.deltanet_chunk_size = 32
# Prefill sub-chunk: controls peak memory per DeltaNet forward call.
# CCCL target = cc_to_min_bytes_in_flight(cc): BI-V100 ≈ lower tier.
# Qwen3.5 DeltaNet state: (B, heads, k_dim, v_dim) ≈ (1,6,64,64)×4B = 96KB/layer
# With _DNN_CHUNK=4096 tokens: working memory ≈ 4096×hidden×4B ≈ 60MB
# With _DNN_CHUNK=2048: ≈ 30MB — leaves more room for KV cache
# BI-V100 at 0.95 GPU util with 256K context needs memory headroom
self.deltanet_prefill_chunk = 4096
# Probe MoE native kernels (CCCL: check op availability per CC)
try:
import ixformer.functions as ixf_F