588-line vllm model implementation based on qwen3_moe.py.
Bootstrap strategy: treat ALL layers as full attention (ignoring
linear_attention optimization). Correct but suboptimal.
Key adaptations:
- _get_text_config(): unwrap composite config -> text_config
- Shared expert support (shared_expert_intermediate_size)
- Skip linear attention weights (conv1d, delta_net, gated_delta)
- Skip vision encoder and MTP weights
- QK norm (Qwen3 style)
- Partial rotary embedding (rope_pct=0.25)
Includes deploy.sh and run_baseline.sh for server deployment.
vllm 0.6.3 KeyError on qwen3_5_moe model type.
Model is hybrid linear+full attention MoE with 256 experts (top-8).
enginex-vllm-bi100-qwen36-main.zip in repo likely contains the fix.
Previous version used base topk policy's bits (11 for key>=2B),
causing SMEM overflow: 512*4*key_size + 2048*4*batches > 49152.
Fix: force bits=8 (same as radix_sort decision for BI-V100).
SMEM: 512*4*key_size + 256*4*batches = manageable.
Also adds while-loop SMEM check on max_batches.
Detected by test_smem_safety.py: 3 overflows at key_size=2,4,8.
Registers all 26 CUB algorithms with metadata:
- 6 'injection' mode: have VLLM_INJECTION_POINTS (reduce/scan/topk/transform/batch_memcpy/for)
- 20 'library' mode: used via CCCL device API, no direct #define injection
- struct_mode: 'named' (bi100_* structs) vs 'inline' (policy_selector returns)
Also adds coverage reporting to generate_patches().
The previous version had a `portioned_smem_per_warp` field that doesn't
exist in CCCL. The actual CCCL RadixSortOnesweepPolicy has:
threads, items, store_algorithm, rank_algorithm, scan_algorithm,
rank_private_partitions, radix_bits
Also adds proper SMEM calculation:
total = max(keys_tile, values_tile, rank_smem) + offsets
with 2KB headroom for kernel stack/locals.
rank_private_partitions set to 1 to minimize SMEM pressure.
Replaces hand-written reduce_threads=512, reduce_items=16 with
_read_reduce_config(accum_size) that reads from tuning_reduce.cuh
via gen_patch.extract_bi100_structs().
Architecture change:
OLD: hand-write values in Python + verify_against_headers() asserts equal
NEW: _read_reduce_config() reads from C++ header (single source of truth)
Falls back to compiled-in defaults only when headers not on disk
(deployed container), with RuntimeWarning.
No hand-written tuning values remain in the normal code path.
verify_against_headers() removed — there is nothing to verify
when there is only one copy of the truth.
Defensive guard: if SMEM cap computes max_threads_by_smem < 32
(or rounds to 0), floor at 32 (one warp). Prevents launching
0 threads which is undefined behavior.
Adds verification that hand-written values in muh_dispatch.py
(reduce_threads=512, reduce_items=16, etc.) match the C++ headers
(bi100_float32_plus_o4 in tuning_reduce.cuh).
Previously: muh_dispatch.py had hand-coded values with no link to
the C++ source of truth. gen_patch.py reads from C++ headers,
but muh_dispatch.py was a separate copy that could diverge.
Now: verify_against_headers() calls gen_patch.extract_bi100_structs()
and compares. Self-test prints mismatches if any exist.
scale_mem_bound now returns {items, threads} (items-first) to match
CCCL's scaling_result struct. All 7 call sites in this file updated.
Previously: auto [t, i] bound threads→t, items→i
Now: auto [i, t] binds items→i, threads→t
The ReducePassPolicy{t, i, ...} constructors remain correct because
they take (threads, items, ...) — t is threads, i is items in both cases.
The old code worked by accident (two reversals canceling out).
1. Return order: {threads, items} → {items, threads} matching CCCL scaling_result
2. Upper clamp: nominal*1 → nominal*2 (CCCL allows small types to double items)
3. Add threads SMEM cap: min(nominal, round_up(48KB/(ts*items), 32))
Verified against all 8 test vectors from CCCL catch2_test_util_arch.cu.
The old code was only safe because current bi100_* structs don't hit the
edge cases — but any future CCCL code copy would silently produce wrong
values.
1. Return order: (items, threads) not (threads, items) — matches CCCL scaling_result
2. Items clamp upper bound: nominal*2, not nominal*1 — allows small types to double
3. Threads SMEM cap: min(nominal, round_up(max_smem/(type*items), 32)) — prevents SMEM overflow
Verified against all 18 CCCL test cases in catch2_test_util_arch.cu (was 4/14, now 18/18).
Note: C++ tuning headers (tuning_reduce.cuh etc.) have corresponding auto [t, i] destructuring
that also needs to flip to auto [i, t]. The bi100_* struct values themselves are correct
(hand-derived from SMEM constraints), but the policy_selector callers of scale_mem_bound
will produce wrong destructuring. Tracked in project/6 as separate fix item.
V1 paged_attention (decode ≤ 8192):
Fix: head_mapping int→Tensor conversion.
VERIFIED: matches manual attention, max diff < 0.001.
Perf: 0.034ms (256 tok), 0.059ms (1K), 0.169ms (4K), 0.272ms (8K).
V2 paged_attention (decode > 8192):
Native V2 kernel EXISTS (ixf_F.vllm_single_query_cached_kv_attention_v2)
but produces INCORRECT output (diff=1.28 vs V1 on same data).
Using Python V2 fallback (paged_attention_v2_pytorch.py) for now.
The native V2 expects [B,H,bs,d] layout (confirmed) but the output
values don't match even with correct layout conversion.
Prefill (flash_attn_func):
VERIFIED: ixf_F.flash_attn_func(q, k, v, causal=True) works
with head_dim=256 and GQA (num_kv_heads=4).
Patched into xformers.py as first-attempt before _run_sdpa_fallback.
Triton: symlinked /usr/local/lib/ → /usr/local/corex/lib64/ for import.
Hardware testing confirmed:
V1: K=[blocks, kv_heads, head_dim/x, block_size, x] (5D), V=[blocks, kv_heads, head_dim, block_size] (4D) → OK
V2: K=[blocks, kv_heads, block_size, head_dim] (4D), V=[blocks, kv_heads, block_size, head_dim] (4D) → OK
V2 with V1's layout → FAIL (Expected key_cache.dim()==4, value_cache.size(3)==head_size)
V1 and V2 use DIFFERENT cache memory layouts in ixformer.
V2 patch now converts cache on the fly before calling native kernel:
K: permute(0,1,3,2,4).reshape → [B,H,bs,d]
V: permute(0,1,3,2).contiguous → [B,H,bs,d]
This is a view+reshape for K (no copy if contiguous) and a transpose+contiguous for V.
The cost is one V copy per decode step, but this enables the native compiled V2 kernel
which is 10-100x faster than the Python fallback it replaces.
Architecture document: docs/paged_attention_kernel_architecture.md
Defines every module from CCCL algorithm patterns before code.
Three-level decomposition from CCCL:
Level 1 (warp_reduce_shfl): shfl.down butterfly for per-thread QK scores
Level 2 (block_reduce_warp_reductions): warp partials → SMEM → block aggregate
Level 3 (agent_scan decoupled lookback): cross-partition combine
Compound type (from summary_statistics.cu):
attention_partial = (max_score, exp_sum, weighted_v[256])
combine(a, b) = online softmax rescaling (same math as Flash Attention)
Key design change: Grid on num_kv_heads, not num_heads.
Before: grid = (1, 24, 200) = 4800 blocks, KV loaded 6x redundantly
After: grid = (1, 4, 200) = 800 blocks, KV loaded once per kv_head
Each block computes GQA_RATIO=6 query heads with shared KV loads.
Reduces KV cache bandwidth by 6x (the GQA ratio).
SMEM budget verified:
K tile [32, 256] fp16 = 16KB
V tile [32, 256] fp16 = 16KB
Total = 32KB ≤ 48KB ✓
Phase 1 kernel: _partition_attn_kernel
Processes query heads sequentially within the GQA group
to minimize register pressure (6 × 256 = 1536 registers
too many if all loaded simultaneously).
Phase 2 kernel: _reduce_partitions_kernel
Also gridded on kv_heads, reduces all partitions for
GQA_RATIO heads per block.
This replaces the previous Triton V2 which was gridded on num_heads
and had no GQA awareness at the kernel level.
Bug: After GQA broadcast optimization, v_perm was [kv_h, seq_len, d]
in the GQA path, but unconditional v_padded allocation used num_heads:
v_padded = torch.zeros((num_heads, padded_len, head_size))
v_padded[:, :seq_len, :] = v_perm # [24, padded, d] vs [4, seq, d] → CRASH
Fix: v_padded/v_parts allocation is now inside the non-GQA else branch.
GQA branch uses its own v_padded_kv with correct [kv_h, padded, d] shape.
This was a real runtime bug — V2 would have crashed on first call
for any GQA model (Qwen3.6, Llama, etc.).
Bug: if l_i > 0 branch in Triton is invalid (compiled as constexpr).
Also: p = exp(scores - m_i_new) computed after m_i_new update was
using the wrong reference max (should subtract m_ij first, then rescale).
Fix: Adapted exactly from prefix_prefill.py's proven-correct pattern:
p = exp(scores - m_ij) # probs relative to chunk max
l_ij = sum(p) # chunk sum
m_i_new = max(m_i, m_ij) # new running max
alpha = exp(m_i - m_i_new) # old accumulator rescale
beta = exp(m_ij - m_i_new) # new chunk rescale
l_i_new = alpha*l_i + beta*l_ij
acc = acc*(alpha*l_i/l_i_new) + (p*beta/l_i_new) @ V
This is the Flash Attention online softmax tiling algorithm.
Same math as CCCL's parallel_reduce with compound accumulators.
Previous commit broadcast Q@K^T (saved 1GB/step).
This commit broadcasts scores@V too (saves 2GB/step).
Before: V expanded from [kv_h, padded_len, d] to [H, padded_len, d]
4×100K×256×4B → 24×100K×256×4B = 400MB → 2.4GB allocation
After: broadcast matmul at kv_h level
se: [kv_h, gqa, P, 1, part_sz] @ V: [kv_h, 1, P, part_sz, d]
→ [kv_h, gqa, P, 1, d] → reshape to [H, P, d]
V stays at kv_h size: 400MB (no 2.4GB allocation)
Total per-decode-step memory for 100K context:
Before all GQA opts: 3.6GB (K expansion + V expansion)
After: 600MB (6x total reduction from GQA ratio=6)
This is the CCCL insight applied: transform_reduce with a compound type.
Instead of expanding to full head count then reducing, keep the reduction
at the minimal group size and broadcast the grouping dimension.
Analysis:
CUDA graph eliminates kernel launch overhead (~10-20% for decode).
At 32768, sequences >32K skip graph capture.
At 65536, most competition workload sequences get graph acceleration.
Memory: CUDA graph capture allocates one copy of all intermediate tensors
at the max captured batch size. With max-num-seqs=1, this is one sequence's
worth of tensors — small relative to model weights.
Combined with V2 enabled for seq>8192 and threshold raised to 65536,
the decode path is now:
seq <= 8192: V1 compiled kernel (fastest)
8192 < seq <= 65536: V2 pytorch (single-bmm, good)
seq > 65536: PyTorch fallback (rare at competition workload)
