Feat/block extend review fixes

.github/workflows/pr-test.yml

@@ -611,6 +611,91 @@ jobs:
       - name: Run H100 Kernel Tests
         run: bash ci/bash.sh ${DOCKER_IMAGE} ./scripts/task_run_unit_tests.sh
 
+  # ---------------------------------------------------------------------------
+  # DLLM Block Extend Attention Tests
+  # ---------------------------------------------------------------------------
+  gpu-tests-dllm-a10g:
+    name: DLLM Block Extend Tests (A10G)
+    needs: [gate, setup]
+    if: |
+      needs.gate.outputs.authorized == 'true' &&
+      needs.setup.outputs.skip_build != 'true' &&
+      github.event.inputs.skip_gpu != 'true'
+    runs-on: [self-hosted, linux, x64, gpu, sm86, spot]
+    timeout-minutes: 60
+    env:
+      DOCKER_IMAGE: flashinfer/flashinfer-ci-cu129:${{ needs.setup.outputs.docker_tag }}
+    steps:
+      - name: Cleanup
+        run: |
+          docker stop $(docker ps -q) 2>/dev/null || true
+          docker rm $(docker ps -aq) 2>/dev/null || true
+          sudo rm -rf ${{ github.workspace }}/* || true
+          sudo rm -rf ${{ github.workspace }}/.[!.]* || true
+          rm -rf ~/.cache/flashinfer_jit || true
+          docker image prune -f || true
+          docker builder prune -f --filter "until=24h" || true
+          nvidia-smi || true
+
+      - uses: actions/checkout@v4
+        with:
+          submodules: recursive
+
+      - name: Start spot termination monitor
+        run: ./scripts/task_monitor_spot.sh &
+
+      - name: Show Node Info
+        run: ./scripts/task_show_node_info.sh
+        env:
+          NODE_NAME: ${{ runner.name }}
+          WORKSPACE: ${{ github.workspace }}
+          BUILD_NUMBER: ${{ github.run_number }}
+
+      - name: Run DLLM Block Extend Tests
+        run: bash ci/bash.sh ${DOCKER_IMAGE} ./scripts/task_jit_run_tests_dllm.sh
+
+  gpu-tests-dllm-h100:
+    name: DLLM Block Extend Tests (H100)
+    needs: [gate, setup]
+    if: |
+      needs.gate.outputs.authorized == 'true' &&
+      needs.setup.outputs.skip_build != 'true' &&
+      github.event.inputs.skip_gpu != 'true'
+    runs-on: [self-hosted, linux, x64, gpu, h100]
+    timeout-minutes: 60
+    env:
+      DOCKER_IMAGE: flashinfer/flashinfer-ci-cu129:${{ needs.setup.outputs.docker_tag }}
+    steps:
+      - name: Cleanup
+        run: |
+          for cid in $(docker ps -q); do
+            if docker inspect "$cid" --format '{{json .Mounts}}' 2>/dev/null | grep -qF "${{ github.workspace }}/"; then
+              docker stop "$cid" 2>/dev/null || true
+            fi
+          done
+          sudo rm -rf ${{ github.workspace }}/* || true
+          sudo rm -rf ${{ github.workspace }}/.[!.]* || true
+          rm -rf ~/.cache/flashinfer_jit || true
+          docker container prune -f 2>/dev/null || true
+          docker image prune -f || true
+          docker builder prune -f --filter "until=24h" || true
+          echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES"
+          nvidia-smi -i "$CUDA_VISIBLE_DEVICES" || true
+
+      - uses: actions/checkout@v4
+        with:
+          submodules: recursive
+
+      - name: Show Node Info
+        run: ./scripts/task_show_node_info.sh
+        env:
+          NODE_NAME: ${{ runner.name }}
+          WORKSPACE: ${{ github.workspace }}
+          BUILD_NUMBER: ${{ github.run_number }}
+
+      - name: Run DLLM Block Extend Tests
+        run: bash ci/bash.sh ${DOCKER_IMAGE} ./scripts/task_jit_run_tests_dllm.sh
+
   # ---------------------------------------------------------------------------
   # Test Results Summary
   # ---------------------------------------------------------------------------
@@ -630,6 +715,8 @@ jobs:
       - analyze-gpu-t4-failure
       - gpu-tests-t4-rerun
       - gpu-tests-h100
+      - gpu-tests-dllm-a10g
+      - gpu-tests-dllm-h100
     runs-on: ubuntu-latest
     steps:
       - name: Check Authorization
@@ -737,6 +824,21 @@ jobs:
             FAILED=true
           fi
 
+          # DLLM tests (no rerun logic - short-running jobs)
+          echo "" >> $GITHUB_STEP_SUMMARY
+          DLLM_A10G="${{ needs.gpu-tests-dllm-a10g.result }}"
+          echo "DLLM Tests (A10G): $DLLM_A10G" >> $GITHUB_STEP_SUMMARY
+          if [ "$DLLM_A10G" != "success" ] && [ "$DLLM_A10G" != "skipped" ] && [ "${{ github.event.inputs.skip_gpu }}" != "true" ]; then
+            FAILED=true
+          fi
+
+          echo "" >> $GITHUB_STEP_SUMMARY
+          DLLM_H100="${{ needs.gpu-tests-dllm-h100.result }}"
+          echo "DLLM Tests (H100): $DLLM_H100" >> $GITHUB_STEP_SUMMARY
+          if [ "$DLLM_H100" != "success" ] && [ "$DLLM_H100" != "skipped" ] && [ "${{ github.event.inputs.skip_gpu }}" != "true" ]; then
+            FAILED=true
+          fi
+
           echo "" >> $GITHUB_STEP_SUMMARY
           if [ "$FAILED" == "true" ]; then
             echo "Result: Tests Failed" >> $GITHUB_STEP_SUMMARY

csrc/batch_prefill_customize_config.jinja

@@ -10,12 +10,7 @@
 #define ADDITIONAL_FUNC_PARAMS {{ additional_func_params }}
 #define ADDITIONAL_PARAMS_SETTER {{ additional_params_setter }}
 
-#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, RaggedParams, PagedParams, ...) \
-  DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \
-    constexpr auto use_custom_mask = MASK_MODE == MaskMode::kCustom; \
-    using AttentionVariant = {{ variant_name }}; \
-    __VA_ARGS__(); \
-  })
+{{ dispatch_context }}
 
 using namespace flashinfer;
 
@@ -76,6 +71,7 @@ struct RaggedParams {
   __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const {
     return kv_indptr[batch_idx + 1] - kv_indptr[batch_idx];
   }
+
 };
 
 struct PagedParams {
@@ -116,6 +112,7 @@ struct PagedParams {
   __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const {
     return paged_kv.get_length(batch_idx);
   }
+
 };
 
 {{ variant_decl }}

csrc/batch_prefill_sm90_customize_config.jinja

@@ -5,11 +5,11 @@
 #include <flashinfer/layout.cuh>
 #include <flashinfer/cutlass_utils.cuh>
 
+#include <flashinfer/utils.cuh>
 #define ADDITIONAL_FUNC_PARAMS {{ additional_func_params }}
 #define ADDITIONAL_PARAMS_SETTER {{ additional_params_setter }}
 
-#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, RaggedParams, PagedParams, ...) \
-  DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = {{ variant_name }}; __VA_ARGS__();})
+{{ dispatch_context }}
 
 using namespace flashinfer;
 
@@ -66,6 +66,16 @@ struct RaggedParams {
   int window_left;
 
   bool causal;
+
+  // Block Expanding support
+  __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const {
+    return qo_lens[batch_idx];
+  }
+
+  __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const {
+    return kv_lens[batch_idx];
+  }
+
 };
 
 struct PagedParams {
@@ -117,6 +127,16 @@ struct PagedParams {
   int window_left;
 
   bool causal;
+
+  // Block Expanding support
+  __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const {
+    return qo_lens[batch_idx];
+  }
+
+  __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const {
+    return kv_lens[batch_idx];
+  }
+
 };
 
 {{ variant_decl }}

csrc/single_prefill_customize_config.jinja

@@ -10,12 +10,7 @@
 #define ADDITIONAL_PARAMS_SETTER {{ additional_params_setter }}
 
 
-#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, POS_ENCODING_MODE, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, USE_FP16_QK_REDUCTION, AttentionVariant, Params, ...) \
-  DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { \
-    constexpr bool use_custom_mask = MASK_MODE == MaskMode::kCustom; \
-    using AttentionVariant = {{ variant_name }}; \
-    __VA_ARGS__(); \
-  })
+{{ dispatch_context }}
 
 
 using namespace flashinfer;

csrc/single_prefill_sm90_customize_config.jinja

@@ -11,8 +11,7 @@
 #define ADDITIONAL_FUNC_PARAMS {{ additional_func_params }}
 #define ADDITIONAL_PARAMS_SETTER {{ additional_params_setter }}
 
-#define DISPATCH_context(DTypeQ, DTypeKV, DTypeO, IdType, MASK_MODE, HEAD_DIM_QK, HEAD_DIM_VO, USE_SLIDING_WINDOW, USE_LOGITS_SOFT_CAP, AttentionVariant, Params, ...) \
-  DISPATCH_MASK_MODE(mask_mode, MASK_MODE, { using AttentionVariant = {{ variant_name }}; __VA_ARGS__(); })
+{{ dispatch_context }}
 
 using namespace flashinfer;
 
@@ -62,6 +61,15 @@ struct Params {
   int window_left;
 
   bool causal;
+
+  // Block Expanding support
+  __host__ __device__ __forceinline__ uint32_t get_qo_len(uint32_t batch_idx) const {
+    return qo_len;
+  }
+
+  __host__ __device__ __forceinline__ uint32_t get_kv_len(uint32_t batch_idx) const {
+    return kv_len;
+  }
 };
 
 {{ variant_decl }}