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:
dylanyunlon
2026-08-03 10:34:28 +00:00
parent dc9ac0a757
commit 8c1955dc92
6 changed files with 77 additions and 51 deletions

View File

@@ -889,9 +889,6 @@ def get_device_attribute(attribute: int, device: int) -> int:
def get_max_shared_memory_per_block_device_attribute(device: int) -> int:
# muh: CONSERVATIVE — keeping 32KB until confirmed on real BI-V100
# hardware.cuh says 48KB, _custom_ops.py says 32KB. One is wrong.
# Test: launch a kernel requesting 33KB SMEM. If it works → 48KB.
return 32 * 1024

View File

@@ -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(

View File

@@ -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'],
)

View File

@@ -338,23 +338,17 @@ def get_default_config(
dtype: Optional[str],
is_marlin: bool,
) -> Dict[str, int]:
# muh: BI-V100 (SM=16, 48KB SMEM) aware defaults
# Qwen3.6 MoE: E≈128, topk=8, K≈2048, N≈5504
# SM=16 → fewer CTAs → each CTA should do more work → larger K tile
# SMEM check: M=64 * K=64 * 2B(fp16) * 2(A+B) = 16KB < 48KB ✓
config = {
'BLOCK_SIZE_M': 64,
'BLOCK_SIZE_N': 64,
'BLOCK_SIZE_K': 64, # muh: 32→64, better memory coalescing on BI-V100
'BLOCK_SIZE_K': 32,
'GROUP_SIZE_M': 8
}
# A heuristic: fused marlin works faster with this config for small M
if M <= E or (is_marlin and M <= 32):
# muh: decode path (M=1 for single-token, M=8 for topk=8)
# BI-V100: K=64 good for memory BW, N=64 for output tile
config = {
'BLOCK_SIZE_M': 16,
'BLOCK_SIZE_N': 64, # muh: 32→64, wider output tile
'BLOCK_SIZE_N': 32,
'BLOCK_SIZE_K': 64,
'GROUP_SIZE_M': 1
}