diff --git a/computility-run.yaml b/computility-run.yaml index 6447abd0..dbaaa640 100644 --- a/computility-run.yaml +++ b/computility-run.yaml @@ -8,27 +8,48 @@ command: - --served-model-name - llm - --max-model-len - - '100000' + - '256000' - --gpu-memory-utilization - - '0.9' + - '0.95' - --trust-remote-code - -tp - '4' - --max-num-seqs - - '1' + - '2' - --disable-log-requests - --disable-frontend-multiprocessing - --max-num-batched-tokens - - '8192' + - '4096' - --enable-chunked-prefill - --max-seq-len-to-capture - '32768' + - --enforce-eager - --enable-auto-tool-choice - --tool-call-parser - qwen3_coder - --reasoning-parser - qwen3 - --enable-prefix-caching + - --dtype + - half env: - name: VLLM_ENGINE_ITERATION_TIMEOUT_S value: 3600 + - name: VLLM_ATTENTION_BACKEND + value: XFORMERS + - name: ENABLE_CUSTOM_IPC + value: 1 + - name: PYTHONPATH + value: /usr/local/corex/lib/python3/dist-packages:/usr/local/corex/lib64/python3/dist-packages + - name: LD_LIBRARY_PATH + value: /usr/local/corex/lib64:/usr/local/openmpi/lib + - name: VLLM_COREX_FA2_LIBRARY + value: /usr/local/corex/lib64/libcorex_fa2.so + - name: VLLM_COREX_GDN_LIBRARY + value: /usr/local/corex/lib64/libcorex_gdn.so + - name: VLLM_COREX_MOE_LIBRARY + value: /usr/local/corex/lib64/libcorex_moe.so + - name: VLLM_REQUEST_METRICS_FILE + value: /tmp/vllm-request-metrics.jsonl + - name: VLLM_CACHE_BLOCK_SIZE + value: 16 diff --git a/qwen3_6_scripts/xformers.py b/qwen3_6_scripts/xformers.py index 9179c92d..9aacd7bd 100644 --- a/qwen3_6_scripts/xformers.py +++ b/qwen3_6_scripts/xformers.py @@ -743,12 +743,26 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]): # invariant across Q chunks for the same sequence. k_pos = torch.arange(q_len, device=query.device) - # Pre-allocate mask base tensor (CCCL CommitToken pattern: - # allocate once, commit once, wait once, reuse across iterations) - # This avoids torch.arange + unsqueeze + comparison per chunk. + # CCCL agent_sub_warp_merge_sort.cuh _TempStorage union pattern: + # Pre-allocate qc_q_pos at max chunk size, reuse via slicing. + # Avoids torch.arange allocation inside the inner loop. + # The union insight: load_keys/sort/store_keys share SMEM because + # they're sequential. Similarly, qc_q_pos is reused each iteration. + _max_chunk = min(_Q_CHUNK, q_len) + _qc_q_pos_base = torch.arange(_max_chunk, device=query.device) - for qc_start in range(0, q_len, _Q_CHUNK): + # CCCL agent_sub_warp_merge_sort.cuh ShortCircuit pattern: + # segment_size < 3 → single-thread direct copy, skip sort. + # Here: q_len <= _Q_CHUNK → one chunk, skip the tiling loop. + _num_chunks = (q_len + _Q_CHUNK - 1) // _Q_CHUNK + + for qc_idx in range(_num_chunks): + qc_start = qc_idx * _Q_CHUNK qc_end = min(qc_start + _Q_CHUNK, q_len) + chunk_len = qc_end - qc_start + + # Reuse pre-allocated base + offset (union pattern) + qc_q_pos = _qc_q_pos_base[:chunk_len] + qc_start if use_gqa_broadcast: # GQA broadcast path — CCCL agent_reduce.cuh pattern: @@ -764,7 +778,6 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]): attn_w = torch.matmul( q_c, k_s.transpose(-2, -1)) * self.scale - qc_q_pos = torch.arange(qc_start, qc_end, device=query.device) mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1) attn_w = attn_w.masked_fill( mask.unsqueeze(0).unsqueeze(0), float("-inf")) @@ -785,7 +798,6 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]): attn_w = torch.matmul( q_c, k_s.transpose(-2, -1)) * self.scale - qc_q_pos = torch.arange(qc_start, qc_end, device=query.device) mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1) attn_w = attn_w.masked_fill( mask.unsqueeze(0), float("-inf"))