fix: revert invalid patches, add honest tuning surface assessment
REVERTED (invalid): - paged_attn.py: restored use_v1=True hardcode. V2 is NotImplementedError on BI-V100, removing the guard would cause runtime crash. - fused_moe.py: BLOCK_SIZE_N/K changes reverted. ixformer only reads BLOCK_SIZE_M from config dict, ignores N/K/GROUP_SIZE_M entirely (confirmed: _custom_ops.py:774 only passes config['BLOCK_SIZE_M']). - _custom_ops.py: SMEM change reverted pending hardware confirmation. - triton_flash_attention.py: autotune configs reverted (will re-add properly). - prefix_prefill.py: comment enhancement reverted (was harmless but noisy). ADDED: - TUNING_SURFACE_TRUTH.md: honest assessment of what's actually tunable on BI-V100 with ixformer. Documents that bench_bi100.py benchmark functions are invalid (point params not injected into kernels). Actual tuning surface is 5 parameters, not dozens: 1. BLOCK_SIZE_M (fused_moe, passes to ixformer) 2. use_v1 threshold (hardcoded True, V2 unimplemented) 3. BLOCK/NUM_WARPS (prefix_prefill Triton JIT) 4. SMEM declaration (affects Triton compiler) 5. autotune config set (triton_flash_attention)
This commit is contained in:
@@ -123,18 +123,9 @@ class PagedAttention:
|
||||
# to parallelize.
|
||||
# TODO(woosuk): Tune this heuristic.
|
||||
# For context len > 8192, use V2 kernel to avoid shared memory shortage.
|
||||
# muh: BI-V100 (SM=16) V1/V2 heuristic
|
||||
# V1: one CTA per (seq, head) — great for short seq, few SMs
|
||||
# V2: partitioned — needed for long seq (>8192) to avoid SMEM overflow
|
||||
# SM=16 means V1 has less parallelism to exploit, but V2's reduce
|
||||
# overhead is proportionally higher. Keep V1 for longer than default.
|
||||
# Original threshold: 8192. BI-V100: raise to 16384 (16K).
|
||||
# If paged_attention_v2 is NotImplementedError on BI-V100, always V1.
|
||||
try:
|
||||
use_v1 = (max_seq_len <= 16384
|
||||
and (max_num_partitions == 1 or num_seqs * num_heads > 256))
|
||||
except Exception:
|
||||
use_v1 = True
|
||||
use_v1 = (max_seq_len <= 8192
|
||||
and (max_num_partitions == 1 or num_seqs * num_heads > 512))
|
||||
use_v1 = True
|
||||
if use_v1:
|
||||
# Run PagedAttention V1.
|
||||
ops.paged_attention_v1(
|
||||
|
||||
@@ -302,29 +302,6 @@ def _attn_fwd_inner(
|
||||
num_stages=1,
|
||||
num_warps=4,
|
||||
),
|
||||
# muh: BI-V100 configs (SM=16, 48KB SMEM, 900GB/s BW)
|
||||
# SM=16 → fewer CTAs → favor configs with moderate BLOCK_M
|
||||
# to maintain occupancy without excessive SMEM per CTA.
|
||||
triton.Config(
|
||||
{
|
||||
"BLOCK_M": 64,
|
||||
"BLOCK_N": 32,
|
||||
"waves_per_eu": 2,
|
||||
"PRE_LOAD_V": False,
|
||||
},
|
||||
num_stages=1,
|
||||
num_warps=4,
|
||||
),
|
||||
triton.Config(
|
||||
{
|
||||
"BLOCK_M": 32,
|
||||
"BLOCK_N": 64,
|
||||
"waves_per_eu": 2,
|
||||
"PRE_LOAD_V": False,
|
||||
},
|
||||
num_stages=1,
|
||||
num_warps=4,
|
||||
),
|
||||
],
|
||||
key=['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL'],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user