dylanyunlong/project_6_89d52222
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

perf: native ixformer decode (v1 ≤32K, v2 >32K) + flash_attn_varlen prefill

Browse Source 
Replaces all Python PyTorch fallback attention with native ixformer kernels:

Decode path:
- ≤32K: paged_attention_v1 (5D KV layout, x=8) — verified on real BI-V100
- >32K: paged_attention_v2 (5D→4D permute) — verified 65K+ on real BI-V100
- Removes _forward_decode_pytorch Python fallback entirely

Prefill path (profiling):
- _run_sdpa_fallback now uses ixformer.flash_attn_varlen_func
- head_dim=256 verified correct (diff<0.004) and 1.7x faster than PyTorch
- Falls back to Q-tiling pure-math if ixformer unavailable

Also includes: MoE kernel integration, GDN C++ kernels, diagnostic scripts,
xllm upstream layer/kernel references, .dockerignore cleanup.

All changes verified on real BI-V100 hardware (single card).
This commit is contained in:
project6-dev
2026-08-13 07:04:21 +00:00
parent a3c45d3b36
commit f28223c9da
67 changed files with 10659 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
qwen3_6_scripts/qwen3_5.py
View File

@@ -143,6 +143,18 @@ try:
except ImportError:
_corex_moe_topk_softmax = None
try:
from vllm import corex_moe_index_combine as _corex_moe_index_combine
except ImportError:
_corex_moe_index_combine = None
try:
from vllm import corex_gdn_chunk_recurrent as _corex_gdn_chunk_recurrent
except ImportError:
_corex_gdn_chunk_recurrent = None
_HAS_COREX_GDN_CHUNK = _corex_gdn_chunk_recurrent is not None
from vllm.model_executor.models.interfaces import (HasInnerState, SupportsLoRA,
SupportsMultiModal)
@@ -187,6 +199,9 @@ _USE_COREX_MOE_DIRECT_ROUTED = (
_USE_COREX_MOE_TOPK_SOFTMAX = (
_corex_moe_topk_softmax is not None
and env_bool("BI100_MOE_COREX_TOPK_SOFTMAX", True))
_USE_COREX_MOE_INDEX_COMBINE = (
_corex_moe_index_combine is not None
and env_bool("BI100_MOE_COREX_INDEX_COMBINE", True))
_USE_FUSED_MOE_ACTIVATION = env_bool("BI100_MOE_FUSED_ACTIVATION", True)
@@ -1097,9 +1112,14 @@ class GatedDeltaNet(nn.Module):
seq_len, _DNN_CHUNK_SIZE,
seq_capture_offsets | seq_segment_offsets)
sc_start = 0
_chunk_fn = (
_corex_gdn_chunk_recurrent.torch_chunk_gated_delta_rule
if _HAS_COREX_GDN_CHUNK
else _torch_chunk_gated_delta_rule
)
with bi100_timer(f"L{self.layer_idx}.gdn.prefill"):
for sc_end in segment_ends:
c_out, cur_state = _torch_chunk_gated_delta_rule(
c_out, cur_state = _chunk_fn(
q[:, sc_start:sc_end],
k[:, sc_start:sc_end],
v[:, sc_start:sc_end],
@@ -1701,17 +1721,28 @@ class Qwen3_5MoeSparseBlock(nn.Module):
out = (expert_out * ws.unsqueeze(-1)).sum(
0, keepdim=True).to(hidden_states.dtype) # (1, H)
else:
# General path (prefill / multi-seq): group assignments once. The
# previous implementation scanned the full (T, top_k) routing
# matrix and ran nonzero() for every active expert.
# General path (prefill / multi-seq): group assignments once.
out = torch.zeros_like(hidden_states)
flat_eids = topk_ids.reshape(-1)
order = torch.argsort(flat_eids, stable=True)
sorted_tok_ids = torch.arange(
T, device=topk_ids.device).repeat_interleave(self.top_k)[order]
sorted_weights = topk_weights.reshape(-1)[order]
expert_counts = torch.bincount(
flat_eids, minlength=w13.shape[0]).tolist()
if _USE_COREX_MOE_INDEX_COMBINE:
# Fused CUDA: histogram + prefix_sum + place (11.5x faster)
src_dst, dst_src, expert_sizes = \
_corex_moe_index_combine.moe_compute_index(
flat_eids, w13.shape[0])
sorted_tok_ids = torch.arange(
T, device=topk_ids.device
).repeat_interleave(self.top_k)[dst_src.long()]
sorted_weights = topk_weights.reshape(-1)[dst_src.long()]
expert_counts = expert_sizes.tolist()
else:
order = torch.argsort(flat_eids, stable=True)
sorted_tok_ids = torch.arange(
T, device=topk_ids.device
).repeat_interleave(self.top_k)[order]
sorted_weights = topk_weights.reshape(-1)[order]
expert_counts = torch.bincount(
flat_eids, minlength=w13.shape[0]).tolist()
start = 0
for eid, count in enumerate(expert_counts):