fix(stem): IndexError on models with more than 38 layers (Qwen3-32B/235B, DeepSeek-R1)

[hobostay]

Problem

Running Stem sparse-attention prefill on a model with more than 38 transformer layers crashes immediately:

IndexError: list index out of range

This affects flagship models AngelSlim explicitly supports, e.g. Qwen3-32B (64 layers) and Qwen3-235B-A22B (94 layers), as well as DeepSeek-R1/V3 (61 layers).

Root cause

generate_exact_k_schedule() reads the per-layer keep ratio from a fixed-length list:

# angelslim/compressor/sparsity/stem/backends/torch_impl.py:49
_DEFAULT_LAYER_KEEP_RATIOS: list[float] = [1.0, 1.0] + [0.2] * 36   # length 38
...
keep_ratio = _DEFAULT_LAYER_KEEP_RATIOS[layer_idx]

layer_idx is assigned to every transformer layer (stem/patch.py sets self.self_attn.layer_idx = i for all layers), so for any model with > 38 layers the lookup overflows the list.

Fix

Express the default keep-ratio as a function of layer_idx instead of a length-38 list. It reproduces the exact previous values for layers 0..37 (1.0 for the first two layers, 0.2 thereafter) and extends the same rule to deeper layers:

-def _DEFAULT_LAYER_KEEP_RATIOS ... = [1.0, 1.0] + [0.2] * 36
+def _default_layer_keep_ratio(layer_idx: int) -> float:
+    return 1.0 if layer_idx < 2 else 0.2
...
-    keep_ratio = _DEFAULT_LAYER_KEEP_RATIOS[layer_idx]
+    keep_ratio = _default_layer_keep_ratio(layer_idx)

Verification

• python3 -m py_compile passes.
• Value-equivalence check:

  old = [1.0, 1.0] + [0.2] * 36
  assert all(_default_layer_keep_ratio(i) == old[i] for i in range(38))   # identical 0..37
  for i in range(38, 94): assert _default_layer_keep_ratio(i) == 0.2      # no longer raises

• Confirmed _DEFAULT_LAYER_KEEP_RATIOS had no other references (only the definition + this call site).

[yisunlp]

Thanks for fixing this, please pass the pre-commit checks

pip3 install pre-commit black isort flake8; cd AngelSlim; pre-commit install;