ShihaoW commited on
Commit
c32291c
·
verified ·
1 Parent(s): fb9f8b1

Document LA Flash implementation details

Browse files
Files changed (1) hide show
  1. kernel_utils/README.md +25 -0
kernel_utils/README.md CHANGED
@@ -25,6 +25,31 @@ types:
25
  | `0` | Full attention over the listed key segment or packed key segments. |
26
  | `1` | Bottom-right causal attention. |
27
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
28
  ## Runtime Knobs
29
 
30
  | Variable | Default | Meaning |
 
25
  | `0` | Full attention over the listed key segment or packed key segments. |
26
  | `1` | Bottom-right causal attention. |
27
 
28
+ ## How It Works
29
+
30
+ `batch_utils.hybrid_runtime` builds sparse range plans for the text decoder.
31
+ Each plan describes which query token intervals attend to which key/value token
32
+ intervals. `kernel_utils.range_attention` executes those plans with
33
+ FlashAttention instead of materializing dense SDPA masks.
34
+
35
+ The runtime follows three paths:
36
+
37
+ - **Packed simple plans:** when each query range maps to one contiguous
38
+ key/value range, LA Flash flattens the selected ranges, builds FlashAttention
39
+ `cu_seqlens_q` / `cu_seqlens_k`, and calls `flash_attn_varlen_func` directly.
40
+ - **Packed MTP full-window plans:** for hybrid MTP decode, multiple full
41
+ key/value windows for the same query block are concatenated into one packed
42
+ key/value sequence before the FlashAttention call. This keeps the sparse
43
+ memory profile without constructing a `[B,H,Q,K]` attention mask.
44
+ - **Compatible multi-segment plans:** when a query range attends to multiple
45
+ segments that cannot be packed as one sequence, each segment is evaluated with
46
+ FlashAttention and the partial outputs are merged with the standard
47
+ log-sum-exp softmax composition.
48
+
49
+ The output tensor shape and dtype match the decoder attention output expected
50
+ by the model. This path is inference-oriented and depends on FlashAttention's
51
+ forward kernels; it is not a custom autograd training backend.
52
+
53
  ## Runtime Knobs
54
 
55
  | Variable | Default | Meaning |