chore: adding lolcats configs scrc and src

configs/experiment/distill_alpaca_clean_xent0_mse1000_lr1e-2.yaml

@@ -0,0 +1,52 @@
+dataset:
+  name: alpaca_clean
+  dataset_config:
+    name: default
+    path: yahma/alpaca-cleaned
+    chunk_size: 1024  # sequence length for distilling
+    concat_data: true
+    cache_dir: 'data/alpaca'  # Change this to where you want to save
+  pretrained_model_config:  # will be updated based on model_config
+    pretrained_model_name_or_path: 'meta-llama/Meta-Llama-3-8B'  
+    cache_dir: '/scratch/'
+  preprocess_config: null
+
+dataloader:
+  batch_size: 1
+  num_workers: 2
+  drop_last: false
+  pin_memory: true
+
+optimizer:
+  optim: adamw_torch_fused
+  lr: 0.01
+  weight_decay: 0.0
+
+lr_scheduler:
+  lr_scheduler_type: reduce_lr_on_plateau
+  mode: min
+  factor: 0.1
+  patience: 10
+  min_lr: 0.00001
+
+trainer:  # HuggingFace Trainer-like arguments  
+  name: distill_attention_xent_mse
+  reverse_kl: false
+  mse_factor: 1000
+  xent_factor: 0
+  
+  bf16: true
+  train_split: train
+  val_split: validation
+  num_train_epochs: 2
+  gradient_accumulation_steps: 8
+  seed: 42
+  batch_size: 1
+  load_best_model_at_end: true
+  greater_is_better: false
+  metric_for_best_model: distill/eval/loss
+  logging_steps: 100
+  evaluation_strategy: steps
+  max_steps: -1
+  eval_steps: 100
+  max_eval_batches: null

configs/experiment/distill_alpaca_clean_xent1_mse1000_lr1e-2.yaml

@@ -0,0 +1,52 @@
+dataset:
+  name: alpaca_clean
+  dataset_config:
+    name: default
+    path: yahma/alpaca-cleaned
+    chunk_size: 1024  # sequence length for distilling
+    concat_data: true
+    cache_dir: 'data/alpaca'  # Change this to where you want to save
+  pretrained_model_config:  # will be updated based on model_config
+    pretrained_model_name_or_path: 'meta-llama/Meta-Llama-3.1-8B'  
+    cache_dir: '/data_persistent2/sim_data/llama-3_1-8b/'
+  preprocess_config: null
+
+dataloader:
+  batch_size: 1
+  num_workers: 2
+  drop_last: false
+  pin_memory: true
+
+optimizer:
+  optim: adamw_torch_fused
+  lr: 0.01
+  weight_decay: 0.0
+
+lr_scheduler:
+  lr_scheduler_type: reduce_lr_on_plateau
+  mode: min
+  factor: 0.1
+  patience: 10
+  min_lr: 0.00001
+
+trainer:  # HuggingFace Trainer-like arguments  
+  name: distill_attention_xent_mse
+  reverse_kl: false
+  mse_factor: 1000
+  xent_factor: 1
+  
+  bf16: true
+  train_split: train
+  val_split: validation
+  num_train_epochs: 2
+  gradient_accumulation_steps: 8
+  seed: 42
+  batch_size: 1
+  load_best_model_at_end: true
+  greater_is_better: false
+  metric_for_best_model: distill/eval/loss
+  logging_steps: 100
+  evaluation_strategy: steps
+  max_steps: -1
+  eval_steps: 100
+  max_eval_batches: null

configs/experiment/eval_alpaca_clean.yaml

@@ -0,0 +1,56 @@
+dataset:
+  name: alpaca_clean
+  dataset_config:
+    name: alpaca
+    path: yahma/alpaca-cleaned
+    chunk_size: 1024  # sequence length for distilling
+    concat_data: true
+    cache_dir: 'data/alpaca'  # Change this to where you want to save
+  pretrained_model_config:
+    pretrained_model_name_or_path: 'mistralai/Mistral-7B-v0.1'  # will be updated based on model_config
+    cache_dir: '/scratch/'
+  preprocess_config: null
+
+dataloader:
+  batch_size: 1
+  num_workers: 2
+  drop_last: false
+  pin_memory: true
+
+optimizer:
+  optim: adamw_torch_fused
+  lr: 1e-4
+  weight_decay: 0.0
+
+lr_scheduler:
+  lr_scheduler_type: reduce_lr_on_plateau
+  mode: min
+  factor: 0.1
+  patience: 10
+  min_lr: 0.00001
+
+trainer:  # HuggingFace Trainer-like arguments  
+  name: finetune_seq2seq
+  bf16: true
+  train_split: train
+  val_split: test
+  num_train_epochs: 2
+  gradient_accumulation_steps: 8
+  seed: 42
+  batch_size: 1
+  load_best_model_at_end: true
+  greater_is_better: true
+  metric_for_best_model: eval/rouge/geometric_mean
+  logging_steps: 100
+  evaluation_strategy: steps
+  max_steps: -1
+  eval_steps: 100
+  max_eval_batches: null
+
+finetune:
+  method: lora
+  kwargs:
+    r: 8
+    lora_alpha: 16
+    lora_dropout: 0  # 0.05
+    target_modules: ['q_proj', 'k_proj', 'v_proj', 'o_proj']

configs/experiment/finetune_lora_fqkvo_alpaca_clean.yaml

@@ -0,0 +1,58 @@
+dataset:
+  name: alpaca_clean
+  dataset_config:
+    name: default
+    path: yahma/alpaca-cleaned
+    chunk_size: 1024
+    concat_data: true
+    cache_dir: "data/alpaca"
+  pretrained_model_config:
+    pretrained_model_name_or_path: "mistralai/Mistral-7B-v0.1" # will be updated based on model_config
+    cache_dir: "/data_persistent2/sim_data/"
+  preprocess_config: null
+
+dataloader:
+  batch_size: 1
+  num_workers: 2
+  drop_last: false
+  pin_memory: true
+
+optimizer:
+  optim: adamw_torch_fused
+  lr: 1e-4
+  weight_decay: 0.0
+
+lr_scheduler:
+  lr_scheduler_type: reduce_lr_on_plateau
+  mode: min
+  factor: 0.1
+  patience: 10
+  min_lr: 0.00001
+
+trainer: # HuggingFace Trainer-like arguments
+  name: default_lm
+  bf16: true
+  train_split: train
+  val_split: validation
+  num_train_epochs: 2
+  gradient_accumulation_steps: 8
+  seed: 42
+  batch_size: 1
+  load_best_model_at_end: true
+  greater_is_better: false
+  metric_for_best_model: eval/loss # eval/rouge/geometric_mean
+  logging_steps: 100
+  evaluation_strategy: steps
+  max_steps: -1
+  eval_steps: 100
+  max_eval_batches: null
+  num_save_ckpt_steps: 200
+
+finetune:
+  method: lora
+  kwargs:
+    r: 8
+    lora_alpha: 16
+    lora_dropout: 0 # 0.05
+    target_modules: ["q_proj", "k_proj", "v_proj", "o_proj"]
+  trainable_weights: ['feature_map_q.mlp.layer', 'feature_map_k.mlp.layer', 'window_factors']

configs/experiment/finetune_lora_qkvo_alpaca_clean.yaml

@@ -0,0 +1,56 @@
+dataset:
+  name: alpaca_clean
+  dataset_config:
+    name: default
+    path: yahma/alpaca-cleaned
+    chunk_size: 1024
+    concat_data: true
+    cache_dir: "data/alpaca"
+  pretrained_model_config:
+    pretrained_model_name_or_path: "mistralai/Mistral-7B-v0.1" # will be updated based on model_config
+    cache_dir: "/scratch/"
+  preprocess_config: null
+
+dataloader:
+  batch_size: 1
+  num_workers: 2
+  drop_last: false
+  pin_memory: true
+
+optimizer:
+  optim: adamw_torch_fused
+  lr: 1e-4
+  weight_decay: 0.0
+
+lr_scheduler:
+  lr_scheduler_type: reduce_lr_on_plateau
+  mode: min
+  factor: 0.1
+  patience: 10
+  min_lr: 0.00001
+
+trainer: # HuggingFace Trainer-like arguments
+  name: default_lm
+  bf16: true
+  train_split: train
+  val_split: validation
+  num_train_epochs: 2
+  gradient_accumulation_steps: 8
+  seed: 42
+  batch_size: 1
+  load_best_model_at_end: true
+  greater_is_better: false
+  metric_for_best_model: eval/loss # eval/rouge/geometric_mean
+  logging_steps: 100
+  evaluation_strategy: steps
+  max_steps: -1
+  eval_steps: 100
+  max_eval_batches: null
+
+finetune:
+  method: lora
+  kwargs:
+    r: 8
+    lora_alpha: 16
+    lora_dropout: 0 # 0.05
+    target_modules: ["q_proj", "k_proj", "v_proj", "o_proj"]

configs/experiment/no_distill_alpaca_clean.yaml

@@ -0,0 +1,29 @@
+dataset:
+  name: alpaca_clean
+  dataset_config:
+    name: alpaca
+    path: yahma/alpaca-cleaned
+    chunk_size: 1024  # sequence length for distilling
+    concat_data: true
+    cache_dir: 'data/alpaca'  # Change this to where you want to save
+  pretrained_model_config:
+    pretrained_model_name_or_path: 'mistralai/Mistral-7B-v0.1'  # will be updated based on model_config
+    cache_dir: '/scr-ssd/mzhang/models/mistral-v0.1'
+  preprocess_config: null
+
+dataloader:
+  batch_size: 1
+  num_workers: 2
+  drop_last: false
+  pin_memory: true
+
+optimizer:
+  optim: adamw_torch_fused
+  lr: 0.01
+  weight_decay: 0.0
+
+lr_scheduler:
+  lr_scheduler_type: none
+
+trainer:  # HuggingFace Trainer-like arguments  
+  name: null

configs/model/base_llama3_1_8b.yaml

@@ -0,0 +1,15 @@
+name: llama
+model:
+  pretrained_model_name_or_path: 'meta-llama/Meta-Llama-3.1-8B'
+  cache_dir: '/scr-ssd/mzhang/models/llama3'  # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 500000.0
+
+attention:
+  attention_type: softmax

configs/model/base_llama3_8b.yaml

@@ -0,0 +1,15 @@
+name: llama
+model:
+  pretrained_model_name_or_path: 'meta-llama/Meta-Llama-3-8B'
+  cache_dir: '/scr-ssd/mzhang/models/llama3'  # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 500000.0
+
+attention:
+  attention_type: softmax

configs/model/base_mistral_7b.yaml

@@ -0,0 +1,15 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "mistralai/Mistral-7B-v0.1"
+  cache_dir: "/scr-ssd/mzhang/models/mistral-7b-v0.1" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 10000.0
+
+attention:
+  attention_type: softmax

configs/model/chunked_experimental/distill_long_llama3_1_8b_lk_smd_wsw64_fd64_w01.yaml

@@ -0,0 +1,40 @@
+# Experimental config for chunked linear attention 
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3.1-8B"
+  cache_dir: "/scr-ssd/mzhang/models/llama-3_1-8b" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 500000.0
+  rope_scaling:
+    factor: 8.0
+    low_freq_factor: 1.0
+    high_freq_factor: 4.0
+    original_max_position_embeddings: 8192
+    rope_type: llama3
+
+attention:
+  attention_type: lolcats_long_llama_window_sw
+  state_chunk_len: 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/chunked_experimental/distill_long_llama3_1_8b_lk_smd_wtk64_fd64_w01.yaml

@@ -0,0 +1,40 @@
+# Experimental config for chunked linear attention 
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3.1-8B"
+  cache_dir: "/scr-ssd/mzhang/models/llama-3_1-8b" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 500000.0
+  rope_scaling:
+    factor: 8.0
+    low_freq_factor: 1.0
+    high_freq_factor: 4.0
+    original_max_position_embeddings: 8192
+    rope_type: llama3
+
+attention:
+  attention_type: lolcats_long_llama_window_tk
+  state_chunk_len: 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/chunked_experimental/distill_long_llama3_8b_lk_smd_wsw64_fd64_w01.yaml

@@ -0,0 +1,34 @@
+# Experimental config for chunked linear attention 
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3-8B"
+  cache_dir: "/scr-ssd/mzhang/models/llama3" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 500000.0
+
+attention:
+  attention_type: lolcats_long_llama_window_sw
+  state_chunk_len: 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/chunked_experimental/distill_long_llama3_8b_lk_smd_wtk64_fd64_w01.yaml

@@ -0,0 +1,34 @@
+# Experimental config for chunked linear attention 
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3-8B"
+  cache_dir: "/scr-ssd/mzhang/models/llama3" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 500000.0
+
+attention:
+  attention_type: lolcats_long_llama_window_tk
+  state_chunk_len: 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/chunked_experimental/distill_long_mistral_7b_lk_smd_wsw64_fd64_w01.yaml

@@ -0,0 +1,36 @@
+# Experimental config for chunked linear attention 
+name: llama
+model:
+  pretrained_model_name_or_path: "mistralai/Mistral-7B-v0.1"
+  cache_dir: "/scr-ssd/mzhang/models/mistral-7b-v0.1" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2 # eager  # so we can load attention weights
+  rope_theta: 10000.0
+
+attention:
+  attention_type: lolcats_long_llama_window_sw
+  state_chunk_len: 512 # 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  train_window_factor: true
+  train_attention_weights: false
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/chunked_experimental/distill_long_mistral_7b_lk_smd_wtk64_fd64_w01.yaml

@@ -0,0 +1,35 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "mistralai/Mistral-7B-v0.1"
+  cache_dir: "/scr-ssd/mzhang/models/mistral-7b-v0.1" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2 # eager  # so we can load attention weights
+  rope_theta: 10000.0
+
+attention:
+  attention_type: lolcats_long_llama_window_tk
+  state_chunk_len: 512 # 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  train_window_factor: true
+  train_attention_weights: false
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_llama3_1_8b_lk_smd_fd64.yaml

@@ -0,0 +1,35 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3.1-8B"
+  cache_dir: "/scr-ssd/mzhang/models/llama-3_1-8b" # Set this to where you want to save checkpoint weights 
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: eager
+  rope_theta: 500000.0
+  rope_scaling:
+    factor: 8.0
+    low_freq_factor: 1.0
+    high_freq_factor: 4.0
+    original_max_position_embeddings: 8192
+    rope_type: llama3
+
+attention:
+  attention_type: lolcats_llama
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_llama3_1_8b_lk_smd_wsw64_fd64_w01.yaml

@@ -0,0 +1,39 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3.1-8B"
+  cache_dir: "/scratch/" # Set this to where you want to save checkpoint weights 
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: eager
+  rope_theta: 500000.0
+  rope_scaling:
+    factor: 8.0
+    low_freq_factor: 1.0
+    high_freq_factor: 4.0
+    original_max_position_embeddings: 8192
+    rope_type: llama3
+
+attention:
+  attention_type: lolcats_llama_window_sw
+  state_chunk_len: 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_llama3_1_8b_lk_smd_wtk64_fd64_w01.yaml

@@ -0,0 +1,39 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3.1-8B"
+  cache_dir: "/scratch/" # Set this to where you want to save checkpoint weights 
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: eager
+  rope_theta: 500000.0
+  rope_scaling:
+    factor: 8.0
+    low_freq_factor: 1.0
+    high_freq_factor: 4.0
+    original_max_position_embeddings: 8192
+    rope_type: llama3
+
+attention:
+  attention_type: lolcats_llama_window_tk
+  state_chunk_len: 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_llama3_1_8b_lk_t2r.yaml

@@ -0,0 +1,35 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3.1-8B"
+  cache_dir: "/scr-ssd/mzhang/models/llama-3_1-8b" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: eager
+  rope_theta: 500000.0
+  rope_scaling:
+    factor: 8.0
+    low_freq_factor: 1.0
+    high_freq_factor: 4.0
+    original_max_position_embeddings: 8192
+    rope_type: llama3
+
+attention:
+  attention_type: lolcats_llama
+  feature_map: relu
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 128
+    skip_connection: false
+    bias: true
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_llama3_8b_lk_smd_fd64.yaml

@@ -0,0 +1,29 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3-8B"
+  cache_dir: "/scr-ssd/mzhang/models/llama3" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 500000.0
+
+attention:
+  attention_type: lolcats_llama
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_llama3_8b_lk_smd_wsw64_fd64_w01.yaml

@@ -0,0 +1,33 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3-8B"
+  cache_dir: "/scr-ssd/mzhang/models/llama3" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 500000.0
+
+attention:
+  attention_type: lolcats_llama_window_sw
+  state_chunk_len: 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_llama3_8b_lk_smd_wtk64_fd64_w01.yaml

@@ -0,0 +1,33 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3-8B"
+  cache_dir: '/scr-ssd/mzhang/models/llama3'  # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 500000.0
+
+attention:
+  attention_type: lolcats_llama_window_tk
+  state_chunk_len: 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_llama3_8b_lk_t2r.yaml

@@ -0,0 +1,29 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "meta-llama/Meta-Llama-3-8B"
+  cache_dir: "/scr-ssd/mzhang/models/llama3" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2
+  rope_theta: 500000.0
+
+attention:
+  attention_type: lolcats_llama
+  feature_map: relu
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 128
+    skip_connection: false
+    bias: true
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_mistral_7b_lk_smd_fd64.yaml

@@ -0,0 +1,29 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "mistralai/Mistral-7B-v0.1"
+  cache_dir: "/scr-ssd/mzhang/models/mistral-7b-v0.1" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2 # eager  # so we can load attention weights
+  rope_theta: 10000.0
+
+attention:
+  attention_type: lolcats_llama
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_mistral_7b_lk_smd_wsw64_fd64_w01.yaml

@@ -0,0 +1,35 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "mistralai/Mistral-7B-v0.1"
+  cache_dir: "/scr-ssd/mzhang/models/mistral-7b-v0.1" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2 # eager  # so we can load attention weights
+  rope_theta: 10000.0
+
+attention:
+  attention_type: lolcats_llama_window_sw
+  state_chunk_len: 512 # 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  train_window_factor: true
+  train_attention_weights: false
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_mistral_7b_lk_smd_wtk64_fd64_w01.yaml

@@ -0,0 +1,35 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "mistralai/Mistral-7B-v0.1"
+  cache_dir: "/scr-ssd/mzhang/models/mistral-7b-v0.1" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2 # eager  # so we can load attention weights
+  rope_theta: 10000.0
+
+attention:
+  attention_type: lolcats_llama_window_tk
+  state_chunk_len: 512 # 1024
+  window_size: 64
+  affine_attention_factors: false
+  init_window_factor: -2.1972245773362196
+  train_window_factor: true
+  train_attention_weights: false
+  feature_map: softmax_dim
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 64
+    skip_connection: false
+    bias: false
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

configs/model/distill_mistral_7b_lk_t2r.yaml

@@ -0,0 +1,29 @@
+name: llama
+model:
+  pretrained_model_name_or_path: "mistralai/Mistral-7B-v0.1"
+  cache_dir: "/scr-ssd/mzhang/models/mistral-7b-v0.1" # Set this to where you want to save checkpoint weights
+  return_dict: true
+  load_in_8bit: false
+  load_in_4bit: false
+  device_map: auto
+  low_cpu_mem_usage: true
+  torch_dtype: bfloat16
+  attn_implementation: flash_attention_2 # eager  # so we can load attention weights
+  rope_theta: 10000.0
+
+attention:
+  attention_type: lolcats_llama
+  feature_map: relu
+  feature_map_kwargs:
+    eps: 1e-12
+    # mlp: null  # to set
+    fullspace: true
+  layer_idx: null # to set
+  learned_kernel: untied_head_einsum
+  learned_kernel_kwargs:
+    feature_dim: 128
+    skip_connection: false
+    bias: true
+    zero_init: false
+  tie_qk_kernels: false
+  train_qk: false

csrc/__init__.py

@@ -0,0 +1,6 @@
+#
+# Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/
+# Written by Angelos Katharopoulos <angelos.katharopoulos@idiap.ch>,
+# Apoorv Vyas <avyas@idiap.ch>
+#
+from .causal_attention import causal_dot_product

csrc/causal_attention.cpp

@@ -0,0 +1,225 @@
+//
+// Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/
+// Written by Angelos Katharopoulos <angelos.katharopoulos@idiap.ch>,
+// Apoorv Vyas <avyas@idiap.ch>
+//
+
+#include <torch/extension.h>
+
+
+/**
+ * Compute a*b^T and save it into out.
+ *
+ * a \in R^A
+ * b \in R^B
+ */
+inline void vvt_dot(float *a, float *b, float *out, int A, int B) {
+    for (int i=0; i<A; i++) {
+        float * bi = b;
+        for (int j=0; j<B; j++) {
+            *out += (*a) * (*bi);
+            out++;
+            bi++;
+        }
+        a++;
+    }
+}
+
+
+/**
+ * Implement a vector matrix product v*m and save it into out.
+ *
+ * v \in R^A
+ * m \in R^{AxB}
+ */
+inline void vm_dot(float *v, float *m, float *out, int A, int B) {
+    // TODO: Consider removing the zeroing part and assuming out already
+    //       contains 0s
+    for (int i=0; i<B; i++) {
+        out[i] = 0;
+    }
+
+    for (int i=0; i<A; i++) {
+        float *oi = out;
+        for (int j=0; j<B; j++) {
+            *oi += (*v) * (*m);
+            oi++;
+            m++;
+        }
+        v++;
+    }
+}
+
+
+/**
+ * Implement a vector transposed-matrix product and save it into out.
+ *
+ * v \in R^B
+ * m \in R^{AxB}
+ */
+inline void vmt_dot(float *v, float *m, float *out, int A, int B) {
+    for (int i=0; i<A; i++) {
+        float *vi = v;
+        float s = 0;
+        for (int j=0; j<B; j++) {
+            s += (*vi) * (*m);
+            vi++;
+            m++;
+        }
+        // TODO: Should we be aggregating? See the comment on vm_dot.
+        *out = s;
+        out++;
+    }
+}
+
+
+/**
+ * Compute the causally masked dot products of queries, keys and values.
+ *
+ * Basically compute V_j' = (Q_{0:j} * K_{0:j}^T) * V_{0:j} for all j. The
+ * computation is done efficiently by changing the order of the dot products.
+ */
+void causal_dot_product(
+    const torch::Tensor queries,
+    const torch::Tensor keys,
+    const torch::Tensor values,
+    torch::Tensor product
+) {
+    // Extract some shapes
+    int N = queries.size(0);
+    int H = queries.size(1);
+    int L = queries.size(2);
+    int E = queries.size(3);
+    int M = values.size(3);
+
+    // Create accessors for all the arguments
+    auto qa = queries.accessor<float, 4>();
+    auto ka = keys.accessor<float, 4>();
+    auto va = values.accessor<float, 4>();
+    auto pa = product.accessor<float, 4>();
+
+    #pragma omp parallel for collapse(2)
+    for (int n=0; n<N; n++) {
+        for (int h=0; h<H; h++) {
+            auto kv = torch::zeros({E, M}, queries.options());
+            float *kvp = kv.data_ptr<float>();
+            for (int l=0; l<L; l++) {
+                vvt_dot(
+                    &ka[n][h][l][0],
+                    &va[n][h][l][0],
+                    kvp,
+                    E,
+                    M
+                );
+                vm_dot(
+                    &qa[n][h][l][0],
+                    kvp,
+                    &pa[n][h][l][0],
+                    E,
+                    M
+                );
+            }
+        }
+    }
+}
+
+
+/**
+ * Compute the gradients of queries, keys and values given the gradient of the
+ * causal_dot_product output.
+ *
+ * Make sure that everything is computed in O(N D^2) complexity.
+ */
+void causal_dot_backward(
+    const torch::Tensor queries,
+    const torch::Tensor keys,
+    const torch::Tensor values,
+    const torch::Tensor grad_out,
+    torch::Tensor grad_queries,
+    torch::Tensor grad_keys,
+    torch::Tensor grad_values
+) {
+    // Extract some shapes
+    int N = queries.size(0);
+    int H = queries.size(1);
+    int L = queries.size(2);
+    int E = queries.size(3);
+    int M = values.size(3);
+
+    // Create accessors for all the arguments
+    auto qa = queries.accessor<float, 4>();
+    auto ka = keys.accessor<float, 4>();
+    auto va = values.accessor<float, 4>();
+    auto ga = grad_out.accessor<float, 4>();
+    auto gqa = grad_queries.accessor<float, 4>();
+    auto gka = grad_keys.accessor<float, 4>();
+    auto gva = grad_values.accessor<float, 4>();
+
+    #pragma omp parallel for collapse(2)
+    for (int n=0; n<N; n++) {
+        for (int h=0; h<H; h++) {
+            auto kv = torch::zeros({E, M}, queries.options());
+            float *kvp = kv.data_ptr<float>();
+
+            // Compute the gradient wrt the queries
+            for (int l=0; l<L; l++) {
+                vvt_dot(
+                    &ka[n][h][l][0],
+                    &va[n][h][l][0],
+                    kvp,
+                    E,
+                    M
+                );
+                vmt_dot(
+                    &ga[n][h][l][0],
+                    kvp,
+                    &gqa[n][h][l][0],
+                    E,
+                    M
+                );
+            }
+
+            // Compute the gradient wrt the keys and values
+            kv.zero_();
+            for (int l=L-1; l>=0; l--) {
+                vvt_dot(
+                    &qa[n][h][l][0],
+                    &ga[n][h][l][0],
+                    kvp,
+                    E,
+                    M
+                );
+                vmt_dot(
+                    &va[n][h][l][0],
+                    kvp,
+                    &gka[n][h][l][0],
+                    E,
+                    M
+                );
+                vm_dot(
+                    &ka[n][h][l][0],
+                    kvp,
+                    &gva[n][h][l][0],
+                    E,
+                    M
+                );
+            }
+        }
+    }
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def(
+        "causal_dot_product",
+        &causal_dot_product,
+        "Compute the weighted sum of values but attending only to previous "
+        "values."
+    );
+    m.def(
+        "causal_dot_backward",
+        &causal_dot_backward,
+        "Compute the gradient of queries, keys and values given the gradient "
+        "of causal_dot_product."
+    );
+}

csrc/causal_attention.py

@@ -0,0 +1,77 @@
+#
+# Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/
+# Written by Angelos Katharopoulos <angelos.katharopoulos@idiap.ch>,
+# Apoorv Vyas <avyas@idiap.ch>
+#
+
+import torch
+
+try:
+    from causal_attention_cuda import causal_dot_product as causal_dot_product_cuda
+    from causal_attention_cuda import causal_dot_backward as causal_dot_backward_cuda
+except ImportError as e:
+    print(e)
+    causal_dot_product_cuda = causal_dot_backward_cuda = None
+
+
+class CausalDotProduct(torch.autograd.Function):
+    """Compute the weighted sum of values but attending only to previous
+    values."""
+    dot = {
+        # "cpu": causal_dot_product_cpu,
+        "cuda": causal_dot_product_cuda
+    }
+    dot_backward = {
+        # "cpu": causal_dot_backward_cpu,
+        "cuda": causal_dot_backward_cuda
+    }
+
+    @staticmethod
+    def forward(ctx, Q, K, V):
+        # Save the inputs for the gradient computation
+        ctx.save_for_backward(Q, K, V)
+
+        # Create the output tensor
+        device = Q.device
+        N, H, L, _ = Q.shape
+        _, _, _, M = V.shape
+        product = torch.zeros((N, H, L, M), dtype=Q.dtype, device=device)
+
+        # Actually perform the dot product
+        CausalDotProduct.dot[device.type](
+            Q.data,
+            K.data,
+            V.data,
+            product
+        )
+        # breakpoint()
+        # CausalDotProduct.dot[device.type](Q.data, K.data, V.data, product)
+
+        return product
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        # Extract the saved tensors
+        Q, K, V = ctx.saved_tensors
+
+        # Allocate memory for the gradients
+        grad_Q = torch.zeros_like(Q)
+        grad_K = torch.zeros_like(K)
+        grad_V = torch.zeros_like(V)
+
+        # Actually compute the gradients
+        CausalDotProduct.dot_backward[Q.device.type](
+            Q.data,
+            K.data,
+            V.data,
+            grad_out,
+            grad_Q,
+            grad_K,
+            grad_V
+        )
+
+        return grad_Q, grad_K, grad_V
+
+
+# Alias the autograd functions to python style snake case naming
+causal_dot_product = CausalDotProduct.apply

csrc/causal_attention_cuda.cu

@@ -0,0 +1,1483 @@
+//
+// Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/
+// Written by Angelos Katharopoulos <angelos.katharopoulos@idiap.ch>,
+// Apoorv Vyas <avyas@idiap.ch>
+//
+
+//
+// For modifications made inside namespace nvidia (authored by jdemouth):
+//
+// Copyright (c) 2021 NVIDIA CORPORATION. All rights reserved.
+// 
+// Permission is hereby granted, free of charge, to any person obtaining a copy of
+// this software and associated documentation files (the "Software"), to deal in
+// the Software without restriction, including without limitation the rights to
+// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+// the Software, and to permit persons to whom the Software is furnished to do so,
+// subject to the following conditions:
+// 
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+// 
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+//
+
+#include <torch/extension.h>
+#include <assert.h>
+#include <stdio.h>
+
+#define ENABLE_NVIDIA_OPTIMIZATIONS
+
+#ifdef ENABLE_NVIDIA_OPTIMIZATIONS
+namespace nvidia {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+constexpr int THREADS_PER_WARP = 32;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+constexpr int LOW_OCCUPANCY_THRESHOLD = 40; // TODO: Make it HW specific (like 1/2 SMs).
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ __host__ int div_up(int m, int n) {
+  return (m + n-1) / n;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ __host__ int round_up(int m, int n) {
+  return div_up(m, n) * n;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename T >
+struct Lmha_params {
+
+  // The output buffer. Dimensions [B, H, L, M].
+  T *out;
+
+  // The input Qs. Dimensions [B, H, L, E].
+  const T *q;
+  // The input Ks. Dimensions [B, H, L, E].
+  const T *k;
+  // The input Vs. Dimensions [B, H, L, M].
+  const T *v;
+
+  // The different dimensions.
+  int B, L, H, E, M;
+
+  // The strides for the different tensors.
+  int q_stride_B, q_stride_H, q_stride_L;
+  int k_stride_B, k_stride_H, k_stride_L;
+  int v_stride_B, v_stride_H, v_stride_L;
+  int o_stride_B, o_stride_H, o_stride_L;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, bool GO_BACKWARD, int WARPS, int COLS_PER_THREAD = 4 >
+__global__ __launch_bounds__(WARPS * THREADS_PER_WARP)
+void lmha_low_occupancy_kernel(Lmha_params<float> params) {
+
+  // The number of threads per block.
+  constexpr int THREADS_PER_BLOCK = WARPS * THREADS_PER_WARP;
+  // The number of rows per thread.
+  constexpr int ROWS_PER_THREAD = E / THREADS_PER_WARP;
+  // The number of steps per iteration.
+  constexpr int COLS_PER_ITER = WARPS * COLS_PER_THREAD;
+
+  // Make sure E is a multiple of the warp size.
+  static_assert(E % THREADS_PER_WARP == 0, "");
+
+  // Shared memory to store V/O.
+  __shared__ float smem_v[COLS_PER_ITER], smem_o[COLS_PER_ITER];
+  // Shared memory buffer to performance the reductions.
+  __shared__ float smem_reds[E * WARPS]; 
+
+  // The sequence processed by that block.
+  const int bi = blockIdx.z;
+  // The head processed by that block.
+  const int hi = blockIdx.y;
+  // The hidden cell in the V/output buffers.
+  const int vi = blockIdx.x;
+
+  // The linear index of the thread.
+  const int tidx = threadIdx.x;
+
+  // Decompose the block in warp/lane.
+  const int warp = tidx / THREADS_PER_WARP;
+  const int lane = tidx % THREADS_PER_WARP;
+
+  // The base offset loaded by the thread in Q and K.
+  int offset_q = bi*params.q_stride_B + hi*params.q_stride_H + lane;
+  int offset_k = bi*params.k_stride_B + hi*params.k_stride_H + lane;
+
+  // If we walk backward, account for the extra offset.
+  if( GO_BACKWARD ) {
+    offset_q += (params.L-1)*params.q_stride_L;
+    offset_k += (params.L-1)*params.k_stride_L;
+  }
+
+  // Position the warp at the beginning of the proper timestep.
+  if( GO_BACKWARD ) {
+    offset_q -= warp*COLS_PER_THREAD*params.q_stride_L;
+    offset_k -= warp*COLS_PER_THREAD*params.k_stride_L;
+  } else {
+    offset_q += warp*COLS_PER_THREAD*params.q_stride_L;
+    offset_k += warp*COLS_PER_THREAD*params.k_stride_L;
+  }
+
+  // Determine the base pointers for Q and K.
+  const float *ptr_q = &params.q[offset_q];
+  const float *ptr_k = &params.k[offset_k];
+
+  // Is a given row valid?
+  int valid_qk[ROWS_PER_THREAD];
+  #pragma unroll
+  for( int ii = 0; ii < ROWS_PER_THREAD; ++ii ) {
+    valid_qk[ii] = lane + ii*THREADS_PER_WARP < params.E;
+  }
+
+  // The offset to the position loaded by the thread in V.
+  int offset_v = bi*params.v_stride_B + hi*params.v_stride_H + vi;
+  int offset_o = bi*params.o_stride_B + hi*params.o_stride_H + vi;
+
+  // If we walk backward, account for the extra offset.
+  if( GO_BACKWARD ) {
+    offset_v += (params.L-1)*params.v_stride_L;
+    offset_o += (params.L-1)*params.o_stride_L;
+  }
+
+  // We load/store a strided matrix of COLS_PER_ITER x OUTPUTS_PER_BLOCK.
+  if( GO_BACKWARD ) {
+    offset_v -= tidx*params.v_stride_L;
+    offset_o -= tidx*params.o_stride_L;
+  } else {
+    offset_v += tidx*params.v_stride_L;
+    offset_o += tidx*params.o_stride_L;
+  }
+
+  // Determine the base pointer for V.
+  const float *ptr_v = &params.v[offset_v];
+  // The output pointer. 
+  float *ptr_o = &params.out[offset_o];
+
+  // The running KVs.
+  float running_kv[ROWS_PER_THREAD];
+  #pragma unroll
+  for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+    running_kv[ri] = 0.f;
+  }
+
+  // Iterate over the timesteps. TODO: Use params.loop_count!!!
+  for( int iter = 0; iter < params.L; iter += COLS_PER_ITER ) {
+
+    // Each thread loads a matrix of elements.
+    float q[ROWS_PER_THREAD][COLS_PER_THREAD], k[ROWS_PER_THREAD][COLS_PER_THREAD];
+
+    // Trigger the memory loads for Q and K.
+    #pragma unroll
+    for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+      #pragma unroll
+      for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+
+        // For Q/K, each warp loads from various timesteps. 
+        int ti = iter + warp*COLS_PER_THREAD;
+        if( GO_BACKWARD ) {
+          ti = params.L - 1 - ti;
+        }
+
+        // Is it a valid access?
+        int valid;
+        if( GO_BACKWARD ) {
+          valid = valid_qk[ri] && ti - ci >= 0;
+        } else {
+          valid = valid_qk[ri] && ti + ci < params.L;
+        }
+
+        // The extra offset to add.
+        if( GO_BACKWARD ) {
+          offset_q = ri*THREADS_PER_WARP - ci*params.q_stride_L;
+          offset_k = ri*THREADS_PER_WARP - ci*params.k_stride_L;
+        } else {
+          offset_q = ri*THREADS_PER_WARP + ci*params.q_stride_L;
+          offset_k = ri*THREADS_PER_WARP + ci*params.k_stride_L;
+        }
+
+        // Load Q/K if they are valid.
+        q[ri][ci] = valid ? ptr_q[offset_q] : 0.f;
+        k[ri][ci] = valid ? ptr_k[offset_k] : 0.f;
+      }
+    }
+
+    // For the V tensor, we assign contiguous thread to different loads. So, ti is different.
+    int ti = iter + tidx;
+    if( GO_BACKWARD ) {
+      ti = params.L - 1 - ti;
+    }
+
+    // Is it a valid access?
+    int valid_vo = tidx < COLS_PER_ITER;
+    if( GO_BACKWARD ) {
+      valid_vo &= ti >= 0;
+    } else {
+      valid_vo &= ti < params.L;
+    }
+
+    // Trigger the loads for V. 
+    float ldg_v = valid_vo ? *ptr_v : 0.f;
+
+    // Move the load pointers.
+    if( GO_BACKWARD ) {
+      ptr_q -= COLS_PER_ITER*params.q_stride_L;
+      ptr_k -= COLS_PER_ITER*params.k_stride_L;
+      ptr_v -= COLS_PER_ITER*params.v_stride_L;
+    } else {
+      ptr_q += COLS_PER_ITER*params.q_stride_L;
+      ptr_k += COLS_PER_ITER*params.k_stride_L;
+      ptr_v += COLS_PER_ITER*params.v_stride_L;
+    }
+
+    // Store to shared memory.
+    if( tidx < COLS_PER_ITER ) {
+      smem_v[tidx] = ldg_v;
+    }
+
+    // Make sure V is in shared memory.
+    __syncthreads();
+
+    // Read V from shared memory.
+    float v[COLS_PER_THREAD];
+    #pragma unroll
+    for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+      v[ci] = smem_v[warp*COLS_PER_THREAD + ci];
+    }
+
+    // Each thread computes local K*V products.
+    float kv[ROWS_PER_THREAD][COLS_PER_THREAD];
+    #pragma unroll
+    for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+      #pragma unroll
+      for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+        kv[ri][ci] = 0.f;
+      }
+    }
+
+    // Update the K*V^T product.
+    #pragma unroll
+    for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+      #pragma unroll
+      for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+        kv[ri][ci] += k[ri][ci] * v[ci];
+      }
+    }
+
+    // We must perform the prefix sums within the thread-block. Start with the thread.
+    #pragma unroll
+    for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+      #pragma unroll
+      for( int ci = 1; ci < COLS_PER_THREAD; ++ci ) {
+        kv[ri][ci] += kv[ri][ci-1];
+      }
+    }
+
+    // Store the partial sums to shared memory. Unless we have no inter-warp reduction to perform.
+    #pragma unroll
+    for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+      smem_reds[warp*E + ri*THREADS_PER_WARP + lane] = kv[ri][COLS_PER_THREAD-1];
+    }
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Each thread deals with one or more column(s) of the matrix.
+    constexpr int SUMS_PER_THREAD = (E + THREADS_PER_BLOCK-1) / THREADS_PER_BLOCK;
+    #pragma unroll
+    for( int ii = 0, idx = tidx; ii < SUMS_PER_THREAD; ++ii, idx += THREADS_PER_BLOCK ) {
+      if( idx < E ) {
+        float sum = smem_reds[idx];
+        #pragma unroll
+        for( int jj = 1; jj < WARPS; ++jj ) {
+          smem_reds[idx + jj*E] = sum += smem_reds[idx + jj*E];
+        }
+      }
+    }
+
+    // Make sure the reductions are stored in shared memory.
+    __syncthreads();
+
+    // Each thread updates his partial products.
+    #pragma unroll
+    for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+      float sum = running_kv[ri];
+      if( warp > 0 ) {
+        sum += smem_reds[(warp-1)*E + lane + ri*THREADS_PER_WARP];
+      }
+      #pragma unroll
+      for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+        kv[ri][ci] += sum;
+      }
+    }
+
+    // Compute the partial output values for that thread.
+    float sum[COLS_PER_THREAD];
+    #pragma unroll
+    for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+      sum[ci] = q[0][ci] * kv[0][ci];
+      #pragma unroll
+      for( int ri = 1; ri < ROWS_PER_THREAD; ++ri ) {
+        sum[ci] += q[ri][ci] * kv[ri][ci];
+      }
+    }
+
+    // Run the parallel reductions inside the warp.
+    #pragma unroll
+    for( int mask = THREADS_PER_WARP / 2; mask >= 1; mask /= 2 ) {
+      #pragma unroll
+      for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+        sum[ci] += __shfl_xor_sync(uint32_t(-1), sum[ci], mask);
+      }
+    }
+
+    // Store the final output to shared memory.
+    if( lane == 0 ) {
+      #pragma unroll
+      for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+        smem_o[warp*COLS_PER_THREAD + ci] = sum[ci];
+      }
+    }
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Store the output.
+    if( valid_vo ) {
+      *ptr_o = smem_o[tidx];
+    }
+
+    // Each thread updates his running kv.
+    #pragma unroll
+    for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+      running_kv[ri] += smem_reds[(WARPS-1)*E + lane + ri*THREADS_PER_WARP];
+    }
+
+    // Move to next location.
+    if( GO_BACKWARD ) {
+      ptr_o -= COLS_PER_ITER*params.o_stride_L;
+    } else {
+      ptr_o += COLS_PER_ITER*params.o_stride_L;
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, bool GO_BACKWARD, int WARPS >
+int lmha_low_occupancy_(const Lmha_params<float> &params) {
+
+  // Make sure we are not going to launch an invalid grid.
+  if( params.H > 65535 || params.B > 65535 ) {
+    return 1;
+  }
+
+  // Prepare the grid and trigger the CUDA kernel.
+  dim3 grid;
+  grid.x = params.M;
+  grid.y = params.H;
+  grid.z = params.B;
+  lmha_low_occupancy_kernel<E, GO_BACKWARD, WARPS><<<grid, WARPS*THREADS_PER_WARP>>>(params);
+  return 0;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, bool GO_BACKWARD >
+int lmha_low_occupancy_(const Lmha_params<float> &params, int blocks) {
+         if( params.M * blocks >= 8*LOW_OCCUPANCY_THRESHOLD ) {
+    return lmha_low_occupancy_<E, GO_BACKWARD,  4>(params);
+  } else if( params.M * blocks >= 4*LOW_OCCUPANCY_THRESHOLD ) {
+    return lmha_low_occupancy_<E, GO_BACKWARD,  8>(params);
+  } else {
+    return lmha_low_occupancy_<E, GO_BACKWARD, 16>(params);
+  }
+  return 1;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, typename Params >
+static inline __device__ __host__ int smem_buffer_elts_(const Params &params) {
+  int M = round_up(params.M, 4);
+  return 2*E + 2*M;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, int THREADS_PER_HEAD, bool GO_BACKWARD >
+__global__ 
+void lmha_kernel(Lmha_params<float> params) {
+
+  // Make sure E is a multiple of 4.
+  static_assert(E % 4 == 0, "");
+
+  // The amount of shared memory per buffer (2 buffers for double-buffering).
+  const int smem_buffer_elts = smem_buffer_elts_<E>(params);
+  // The M dimension for shared memory.
+  const int M = round_up(params.M, 4);
+
+  // Shared memory to store Q, K and V. Size is 2*smem_buffer_elts.
+  extern __shared__ float smem_[];
+
+  // The various shared memory buffers.
+  float *smem_q = &smem_[0*E];
+  float *smem_k = &smem_[1*E];
+  float *smem_v = &smem_[2*E];
+  float *smem_o = &smem_[2*E + M];
+
+  // The index of the shared memory buffer (for double-buffering).
+  int smem_curr = 0;
+
+  // The sequence processed by that block.
+  const int bi = blockIdx.y;
+  // The head processed by that block.
+  const int hi = blockIdx.x;
+
+  // The linear index of the thread.
+  const int tidx = threadIdx.x;
+
+  // The offset to the position loaded by the thread in Q.
+  int offset_q = bi*params.q_stride_B + hi*params.q_stride_H + tidx;
+  // The offset to the position loaded by the thread in K.
+  int offset_k = bi*params.k_stride_B + hi*params.k_stride_H + tidx;
+
+  // If we walk backward, account for the extra offset.
+  if( GO_BACKWARD ) {
+    offset_q += (params.L-1)*params.q_stride_L;
+    offset_k += (params.L-1)*params.k_stride_L;
+  }
+
+  // Determine the base pointers for Q and K.
+  const float *ptr_q = &params.q[offset_q];
+  const float *ptr_k = &params.k[offset_k];
+
+  // The offset to the position loaded by the thread in V and O.
+  int offset_v = bi*params.v_stride_B + hi*params.v_stride_H + tidx;
+  int offset_o = bi*params.o_stride_B + hi*params.o_stride_H + tidx;
+
+  // If we walk backward, account for the extra offset.
+  if( GO_BACKWARD ) {
+    offset_v += (params.L-1)*params.v_stride_L;
+    offset_o += (params.L-1)*params.o_stride_L;
+  }
+
+  // Determine the base pointers for V.
+  const float *ptr_v = &params.v[offset_v];
+
+  // Is it an active Q/K thread?
+  const int active_qk = tidx < params.E;
+
+  // Trigger the memory loads for Q and K.
+  float ldg_q = 0.f, ldg_k = 0.f;
+  if( active_qk ) {
+    ldg_q = *ptr_q;
+    ldg_k = *ptr_k;
+  }
+
+  // Is it an active V thread?
+  const int active_v = tidx < params.M;
+
+  // Trigger the memory loads for V. 
+  float ldg_v = 0.f;
+  if( active_v ) {
+    ldg_v = *ptr_v;
+  }
+
+  // Move the load pointers.
+  if( GO_BACKWARD ) {
+    ptr_q -= params.q_stride_L;
+    ptr_k -= params.k_stride_L;
+    ptr_v -= params.v_stride_L;
+  } else {
+    ptr_q += params.q_stride_L;
+    ptr_k += params.k_stride_L;
+    ptr_v += params.v_stride_L;
+  }
+
+  // The number of FLOAT4s per head.
+  constexpr int FLOAT4s_PER_HEAD = E / 4;
+  // The number of FLOAT4s per thread.
+  constexpr int FLOAT4s_PER_THREAD = FLOAT4s_PER_HEAD / THREADS_PER_HEAD;
+
+  // The storage for the K*V^T values.
+  float4 kv[FLOAT4s_PER_THREAD]; 
+  #pragma unroll
+  for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+    kv[ii] = make_float4(0.f, 0.f, 0.f, 0.f);
+  }
+
+  // The output pointer.
+  float *out_ptr = &params.out[offset_o];
+
+  // Store to shared memory Q and K.
+  if( tidx < E ) { 
+    smem_q[smem_curr*smem_buffer_elts + tidx] = ldg_q; 
+    smem_k[smem_curr*smem_buffer_elts + tidx] = ldg_k; 
+  }
+
+  // Store to shared memory V. All threads store valid values.
+  if( tidx < M ) {
+    smem_v[smem_curr*smem_buffer_elts + tidx] = ldg_v;
+  }
+
+  // The position of the thread in the V dimension.
+  int vo = tidx / THREADS_PER_HEAD;
+  int vi = tidx % THREADS_PER_HEAD;
+
+  // Iterate over the timesteps.
+  for( int ti = 0; ti < params.L; ++ti ) {
+
+    // Is it the last iteration?
+    int is_last = ti == params.L - 1;
+
+    // Trigger the next loads for Q and K.
+    if( !is_last && active_qk ) {
+      ldg_q = *ptr_q;
+      ldg_k = *ptr_k;
+    }
+
+    // Trigger the next loads for V.
+    if( !is_last && active_v ) {
+      ldg_v = *ptr_v;
+    }
+
+    // Move the load pointers.
+    if( GO_BACKWARD ) {
+      ptr_q -= params.q_stride_L;
+      ptr_k -= params.k_stride_L;
+      ptr_v -= params.v_stride_L;
+    } else {
+      ptr_q += params.q_stride_L;
+      ptr_k += params.k_stride_L;
+      ptr_v += params.v_stride_L;
+    }
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Each thread loads 4 values from K.
+    float4 k[FLOAT4s_PER_THREAD];
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      int ki = tidx % THREADS_PER_HEAD * 4 + ii * THREADS_PER_HEAD * 4;
+      k[ii] = *reinterpret_cast<const float4*>(&smem_k[smem_curr*smem_buffer_elts + ki]);
+    }
+
+    // Each thread loads a single V value.
+    float v = 0.f;
+    if( vo < params.M ) {
+      v = *reinterpret_cast<const float *>(&smem_v[smem_curr*smem_buffer_elts + vo]);
+    }
+
+    // Update the K*V^T product.
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      kv[ii].x += k[ii].x * v;
+      kv[ii].y += k[ii].y * v;
+      kv[ii].z += k[ii].z * v;
+      kv[ii].w += k[ii].w * v;
+    }
+
+    // Load the Q values from shared memory.
+    float4 q[FLOAT4s_PER_THREAD]; 
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      int qi = tidx % THREADS_PER_HEAD * 4 + ii * THREADS_PER_HEAD * 4;
+      q[ii] = *reinterpret_cast<const float4*>(&smem_q[smem_curr*smem_buffer_elts + qi]);
+    }
+
+    // Compute the partial output value for that thread.
+    float sum = 0.f;
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      sum += q[ii].x * kv[ii].x;
+      sum += q[ii].y * kv[ii].y;
+      sum += q[ii].z * kv[ii].z;
+      sum += q[ii].w * kv[ii].w;
+    }
+
+    // Finalize the computation of the sum (if we have more than 1 thread per head).
+    if( THREADS_PER_HEAD > 1 ) {
+
+      // Finalize the sum for each head.
+      #pragma unroll
+      for( int mask = THREADS_PER_HEAD / 2; mask >= 1; mask /= 2 ) {
+        sum += __shfl_xor_sync(uint32_t(-1), sum, mask);
+      }
+
+      // Store to shared memory.
+      if( vo < M && vi == 0 ) {
+        smem_o[smem_curr*smem_buffer_elts + vo] = sum;
+      }
+
+      // Make sure the data is in shared memory.
+      __syncthreads();
+
+      // Active threads read the data to store.
+      if( active_v ) {
+        sum = smem_o[smem_curr*smem_buffer_elts + tidx];
+      }
+
+    } // THREADS_PER_HEAD > 1.
+
+    // Store the output. All the threads are active.
+    if( active_v ) {
+      *out_ptr = sum;
+    }
+
+    // Move to next location.
+    if( GO_BACKWARD ) {
+      out_ptr -= params.o_stride_L;
+    } else {
+      out_ptr += params.o_stride_L;
+    }
+
+    // Move the shared memory buffer.
+    smem_curr = (smem_curr + 1) % 2;
+
+    // Store to shared memory for Q and K.
+    if( !is_last && tidx < E ) {
+      smem_q[smem_curr*smem_buffer_elts + tidx] = ldg_q;
+      smem_k[smem_curr*smem_buffer_elts + tidx] = ldg_k;
+    }
+
+    // Store to shared memory for V.
+    if( !is_last && tidx < M ) {
+      smem_v[smem_curr*smem_buffer_elts + tidx] = ldg_v;
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, int THREADS_PER_HEAD, bool GO_BACKWARD >
+int lmha_(const Lmha_params<float> &params) {
+  // The M dimension rounded up to 4.
+  int M = round_up(params.M, 4);
+
+  // The number of threads in the block.
+  int block = round_up(max(E, M*THREADS_PER_HEAD), 32);
+  if( block > 512 || params.B > 65535 ) {
+    return 1;
+  }
+
+  // Prepare the kernel.
+  dim3 grid(params.H, params.B);
+  size_t smem = smem_buffer_elts_<E>(params)*2*sizeof(float);
+  lmha_kernel<E, THREADS_PER_HEAD, GO_BACKWARD><<<grid, block, smem>>>(params);
+  return 0;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< bool GO_BACKWARD >
+int lmha(const Lmha_params<float> &params) {
+  int blocks = params.B * params.H;
+  int res = 1;
+  if( blocks < LOW_OCCUPANCY_THRESHOLD ) { 
+           if( params.E <=  32 ) {
+      res = lmha_low_occupancy_< 32, GO_BACKWARD>(params, blocks);
+    } else if( params.E <=  64 ) {
+      res = lmha_low_occupancy_< 64, GO_BACKWARD>(params, blocks);
+    } else if( params.E <= 128 ) {
+      res = lmha_low_occupancy_<128, GO_BACKWARD>(params, blocks);
+    } else if( params.E <= 256 ) {
+      res = lmha_low_occupancy_<256, GO_BACKWARD>(params, blocks);
+    }
+  } else {
+           if( params.E <=  32 ) {
+      res = lmha_< 32, 1, GO_BACKWARD>(params);
+    } else if( params.E <=  48 ) {
+      res = lmha_< 48, 1, GO_BACKWARD>(params);
+    } else if( params.E <=  64 ) {
+      res = lmha_< 64, 1, GO_BACKWARD>(params);
+    } else if( params.E <= 128 ) {
+      res = lmha_<128, 2, GO_BACKWARD>(params);
+    } else if( params.E <= 256 ) {
+      res = lmha_<256, 4, GO_BACKWARD>(params);
+    }
+  }
+  return res;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename T >
+inline void set_params(Lmha_params<T> &params,
+                       const torch::Tensor q,
+                       const torch::Tensor k,
+                       const torch::Tensor v,
+                       torch::Tensor       o) {
+
+  // Define the pointers.
+  params.out = o.data_ptr<T>();
+  params.q   = q.data_ptr<T>();
+  params.k   = k.data_ptr<T>();
+  params.v   = v.data_ptr<T>();
+
+  // Define the strides.
+  params.q_stride_B = (int) q.stride(0);
+  params.q_stride_H = (int) q.stride(1);
+  params.q_stride_L = (int) q.stride(2);
+  params.k_stride_B = (int) k.stride(0);
+  params.k_stride_H = (int) k.stride(1);
+  params.k_stride_L = (int) k.stride(2);
+  params.v_stride_B = (int) v.stride(0);
+  params.v_stride_H = (int) v.stride(1);
+  params.v_stride_L = (int) v.stride(2);
+  params.o_stride_B = (int) o.stride(0);
+  params.o_stride_H = (int) o.stride(1);
+  params.o_stride_L = (int) o.stride(2);
+
+  // Extract the dimensions.
+  int N = q.size(0);
+  int H = q.size(1);
+  int L = q.size(2);
+  int E = q.size(3);
+  int M = v.size(3);
+
+  params.B = N;
+  params.L = L;
+  params.H  = H;
+  params.E = E;
+  params.M = M;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+int lmha_fwd(const torch::Tensor queries,
+             const torch::Tensor keys,
+             const torch::Tensor values,
+             torch::Tensor product) {
+
+  // Make sure that we are using the correct GPU device
+  torch::DeviceGuard _guard(queries.device());
+
+  // Make sure the inner-most dimension of the tensors is packed.
+  assert(queries.stride(3) == 1);
+  assert(keys   .stride(3) == 1);
+  assert(values .stride(3) == 1);
+  assert(product.stride(3) == 1);
+
+  // Extract the dimensions.
+  int N = queries.size(0);
+  int H = queries.size(1);
+  int L = queries.size(2);
+  int E = queries.size(3);
+  int M = values.size (3);
+
+  // The structure of params.
+  Lmha_params<float> params;
+  set_params(params, queries, keys, values, product);
+
+  // Launch the kernel.
+  return lmha<false>(params);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename T >
+struct Lmha_bwd_params {
+
+  // The output buffer for K. Dimensions [B, H, L, D].
+  T *out_k;
+  // The output buffer for V. Dimensions [B, H, L, D].
+  T *out_v;
+
+  // The input Qs. Dimensions [B, H, L, D].
+  const T *q;
+  // The input Ks. Dimensions [B, H, L, D].
+  const T *k;
+  // The input Vs. Dimensions [B, H, L, D].
+  const T *v;
+  // The input Gs. Dimensions [B, H, L, D].
+  const T *g;
+
+  // The dimensions.
+  int B, L, H, M, E;
+
+  // The strides for the input tensors.
+  int q_stride_B, q_stride_L, q_stride_H;
+  int k_stride_B, k_stride_L, k_stride_H;
+  int v_stride_B, v_stride_L, v_stride_H;
+  int g_stride_B, g_stride_L, g_stride_H;
+
+  // The strides for the outputs.
+  int out_k_stride_B, out_k_stride_L, out_k_stride_H;
+  int out_v_stride_B, out_v_stride_L, out_v_stride_H;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int D, int THREADS_PER_HEAD >
+__global__ __launch_bounds__(D*THREADS_PER_HEAD*2)
+void lmha_bwd_kernel(Lmha_bwd_params<float> params) {
+
+  // Make sure D is a multiple of 4.
+  static_assert(D % 4 == 0, "");
+
+  // The shared memory buffers.
+  __shared__ struct Smem { float qg[2*D], kv[2*D], out_kv[2*D]; } smem_[2];
+
+  // The index of the shared memory buffer (for double-buffering).
+  int smem_curr = 0;
+
+  // The sequence processed by that block.
+  const int bi = blockIdx.y;
+  // The head processed by that block.
+  const int hi = blockIdx.x;
+
+  // The linear index of the thread.
+  const int tidx = threadIdx.x;
+
+  // Split the threads into two slices.
+  int so = tidx / (D*THREADS_PER_HEAD);
+  int si = tidx % (D*THREADS_PER_HEAD);
+
+  // The strides for B/L/H for the Q/G tensors.
+  int qg_stride_B, qg_stride_L, qg_stride_H;
+  if( so == 0 ) {
+    qg_stride_B = params.q_stride_B;
+    qg_stride_L = params.q_stride_L;
+    qg_stride_H = params.q_stride_H;
+  } else {
+    qg_stride_B = params.g_stride_B;
+    qg_stride_L = params.g_stride_L;
+    qg_stride_H = params.g_stride_H;
+  }
+
+  // The strides for B/L/H for the K/V tensors.
+  int kv_stride_B, kv_stride_L, kv_stride_H;
+  if( so == 0 ) {
+    kv_stride_B = params.k_stride_B;
+    kv_stride_L = params.k_stride_L;
+    kv_stride_H = params.k_stride_H;
+  } else {
+    kv_stride_B = params.v_stride_B;
+    kv_stride_L = params.v_stride_L;
+    kv_stride_H = params.v_stride_H;
+  }
+
+  // The hidden size.
+  int hidden_size_per_head = 0;
+  if( so == 0 ) {
+    hidden_size_per_head = params.E;
+  } else {
+    hidden_size_per_head = params.M;
+  }
+
+  // Where to start reading from.
+  int offset_qg = bi*qg_stride_B + hi*qg_stride_H + si;
+  int offset_kv = bi*kv_stride_B + hi*kv_stride_H + si;
+
+  // We walk backward, account for the extra offset.
+  offset_qg += (params.L-1)*qg_stride_L;
+  offset_kv += (params.L-1)*kv_stride_L;
+
+  // Determine the base pointers for Q, K, V and G.
+  const float *ptr_qg = &(so == 0 ? params.q : params.g)[offset_qg];
+  const float *ptr_kv = &(so == 0 ? params.k : params.v)[offset_kv]; 
+
+  // Is it an active thread?
+  const int active = si < hidden_size_per_head;
+
+  // Trigger the memory loads for Q, K, V and G.
+  float ldg_qg = 0.f, ldg_kv = 0.f;
+  if( active ) {
+    ldg_qg = *ptr_qg;
+    ldg_kv = *ptr_kv;
+  }
+
+  // Move the load pointers (backward).
+  ptr_qg -= qg_stride_L;
+  ptr_kv -= kv_stride_L;
+
+  // The number of FLOAT4s per head.
+  constexpr int FLOAT4s_PER_HEAD = D / 4;
+  // The number of FLOAT4s per thread.
+  constexpr int FLOAT4s_PER_THREAD = FLOAT4s_PER_HEAD / THREADS_PER_HEAD;
+
+  // The storage for the G*Q^T or Q^T*G values.
+  float4 gq[FLOAT4s_PER_THREAD]; 
+  #pragma unroll
+  for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+    gq[ii] = make_float4(0.f, 0.f, 0.f, 0.f);
+  }
+
+  // The strides for B/L/H for the K/V tensors.
+  int out_kv_stride_B, out_kv_stride_L, out_kv_stride_H;
+  if( so == 0 ) {
+    out_kv_stride_B = params.out_k_stride_B;
+    out_kv_stride_L = params.out_k_stride_L;
+    out_kv_stride_H = params.out_k_stride_H;
+  } else {
+    out_kv_stride_B = params.out_v_stride_B;
+    out_kv_stride_L = params.out_v_stride_L;
+    out_kv_stride_H = params.out_v_stride_H;
+  }
+
+  // Where to start reading from.
+  int offset_out_kv = bi*out_kv_stride_B + hi*out_kv_stride_H + si;
+
+  // We walk backward, account for the extra offset.
+  offset_out_kv += (params.L-1)*out_kv_stride_L;
+
+  // The output pointer.
+  float *ptr_out_kv = &(so == 0 ? params.out_k : params.out_v)[offset_out_kv];
+
+  // Store to shared memory.
+  if( si < D ) { 
+    smem_[smem_curr].qg[so*D + si] = ldg_qg; 
+    smem_[smem_curr].kv[so*D + si] = ldg_kv; 
+  }
+
+  // The position of the thread in the output dimension.
+  int oo = si / THREADS_PER_HEAD % D;
+  int oi = si % THREADS_PER_HEAD * 4;
+
+  // Iterate over the timesteps.
+  for( int ti = 0; ti < params.L; ++ti ) {
+
+    // Is it the last iteration?
+    int is_last = ti == params.L - 1;
+
+    // Trigger the next loads.
+    if( !is_last && active ) {
+      ldg_qg = *ptr_qg;
+      ldg_kv = *ptr_kv;
+    }
+
+    // Move the load pointers.
+    ptr_qg -= qg_stride_L;
+    ptr_kv -= kv_stride_L;
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Each thread loads 4 values from G or Q.
+    float4 g[FLOAT4s_PER_THREAD];
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      float *smem_ptr = &smem_[smem_curr].qg[(so^1)*D + oi];
+      g[ii] = *reinterpret_cast<const float4*>(&smem_ptr[ii*THREADS_PER_HEAD*4]);
+    }
+
+    // Each thread loads a single from Q or G value.
+    float q = smem_[smem_curr].qg[so*D + oo];
+
+    // Update the G*Q^T or Q*G^T product.
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      gq[ii].x += g[ii].x * q;
+      gq[ii].y += g[ii].y * q;
+      gq[ii].z += g[ii].z * q;
+      gq[ii].w += g[ii].w * q;
+    }
+
+    // Load the V or K values from shared memory.
+    float4 v[FLOAT4s_PER_THREAD]; 
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      float *smem_ptr = &smem_[smem_curr].kv[(so^1)*D + oi];
+      v[ii] = *reinterpret_cast<const float4*>(&smem_ptr[ii*THREADS_PER_HEAD*4]);
+    }
+
+    // Compute the partial output value for that thread.
+    float sum = 0.f;
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      sum += v[ii].x * gq[ii].x;
+      sum += v[ii].y * gq[ii].y;
+      sum += v[ii].z * gq[ii].z;
+      sum += v[ii].w * gq[ii].w;
+    }
+
+    // Finalize the computation of the sum (if we have more than 1 thread per head).
+    if( THREADS_PER_HEAD > 1 ) {
+
+      // Finalize the sum for each head.
+      #pragma unroll
+      for( int mask = THREADS_PER_HEAD / 2; mask >= 1; mask /= 2 ) {
+        sum += __shfl_xor_sync(uint32_t(-1), sum, mask);
+      }
+
+      // Store to shared memory.
+      if( oi == 0 ) {
+        smem_[smem_curr].out_kv[so*D + oo] = sum;
+      }
+
+      // Make sure the data is in shared memory.
+      __syncthreads();
+
+      // Active threads read the data to store.
+      if( si < hidden_size_per_head ) {
+        sum = smem_[smem_curr].out_kv[so*D + si];
+      }
+
+    } // THREADS_PER_HEAD > 1.
+
+    // Store the output. All the threads are active.
+    if( si < hidden_size_per_head ) {
+      *ptr_out_kv = sum;
+    }
+
+    // Move to next location.
+    ptr_out_kv -= out_kv_stride_L;
+
+    // Move the shared memory buffer.
+    smem_curr = (smem_curr + 1) % 2;
+
+    // Store to shared memory for Q and K.
+    if( !is_last && si < D ) {
+      smem_[smem_curr].qg[so*D + si] = ldg_qg; 
+      smem_[smem_curr].kv[so*D + si] = ldg_kv; 
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int D, int THREADS_PER_HEAD >
+int lmha_bwd_(const Lmha_bwd_params<float> &params) {
+  int block = D*THREADS_PER_HEAD*2;
+  if( block >= 1024 || params.B > 65535 ) {
+    return 1;
+  }
+  dim3 grid(params.H, params.B);
+  lmha_bwd_kernel<D, THREADS_PER_HEAD><<<grid, block>>>(params);
+  return 0;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+int lmha_bwd(const Lmha_bwd_params<float> &params) {
+  int blocks = params.B * params.H;
+  if( blocks < LOW_OCCUPANCY_THRESHOLD ) { 
+    return 1;
+  }
+
+  int hidden_size_per_head = max(params.E, params.M);
+  int res = 1;
+  if( hidden_size_per_head <= 32 ) {
+    res = lmha_bwd_< 32, 1>(params);
+  } else if( hidden_size_per_head <= 64 ) {
+    res = lmha_bwd_< 64, 1>(params);
+  } else if( hidden_size_per_head <= 128 ) {
+    res = lmha_bwd_<128, 2>(params);
+  } else if( hidden_size_per_head <= 256 ) {
+    res = lmha_bwd_<256, 4>(params);
+  }
+  return res;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+int lmha_bwd(const torch::Tensor queries,
+             const torch::Tensor keys,
+             const torch::Tensor values,
+             const torch::Tensor grad_out,
+             torch::Tensor grad_queries,
+             torch::Tensor grad_keys,
+             torch::Tensor grad_values) {
+
+  // Make sure that we are using the correct GPU device
+  torch::DeviceGuard _guard(queries.device());
+
+  // Make sure the inner-most dimension of the tensors is packed.
+  assert(queries     .stride(3) == 1);
+  assert(keys        .stride(3) == 1);
+  assert(values      .stride(3) == 1);
+  assert(grad_out    .stride(3) == 1);
+  assert(grad_queries.stride(3) == 1);
+  assert(grad_keys   .stride(3) == 1);
+  assert(grad_values .stride(3) == 1);
+
+  // Extract the dimensions.
+  int N = queries.size(0);
+  int H = queries.size(1);
+  int L = queries.size(2);
+  int E = queries.size(3);
+  int M = values.size (3);
+
+  // Gradient on Q.
+
+  // The structure of params.
+  Lmha_params<float> params;
+  set_params(params, grad_out, values, keys, grad_queries);
+
+  // Launch the kernel.
+  int res = lmha<false>(params);
+  if( res ) {
+    return res;
+  }
+
+  // Gradient on K and V together.
+
+  Lmha_bwd_params<float> bwd_params;
+  bwd_params.out_k = grad_keys.data_ptr<float>();
+  bwd_params.out_v = grad_values.data_ptr<float>();
+  bwd_params.q = queries.data_ptr<float>();
+  bwd_params.k = keys.data_ptr<float>();
+  bwd_params.v = values.data_ptr<float>();
+  bwd_params.g = grad_out.data_ptr<float>();
+
+  bwd_params.B = N;
+  bwd_params.L = L;
+  bwd_params.H = H;
+  bwd_params.E = E;
+  bwd_params.M = M;
+
+  bwd_params.q_stride_B = queries.stride(0);
+  bwd_params.q_stride_H = queries.stride(1);
+  bwd_params.q_stride_L = queries.stride(2);
+  bwd_params.k_stride_B = keys.stride(0);
+  bwd_params.k_stride_H = keys.stride(1);
+  bwd_params.k_stride_L = keys.stride(2);
+  bwd_params.v_stride_B = values.stride(0);
+  bwd_params.v_stride_H = values.stride(1);
+  bwd_params.v_stride_L = values.stride(2);
+  bwd_params.g_stride_B = grad_out.stride(0);
+  bwd_params.g_stride_H = grad_out.stride(1);
+  bwd_params.g_stride_L = grad_out.stride(2);
+
+  bwd_params.out_k_stride_B = grad_keys.stride(0);
+  bwd_params.out_k_stride_H = grad_keys.stride(1);
+  bwd_params.out_k_stride_L = grad_keys.stride(2);
+  bwd_params.out_v_stride_B = grad_values.stride(0);
+  bwd_params.out_v_stride_H = grad_values.stride(1);
+  bwd_params.out_v_stride_L = grad_values.stride(2);
+
+  // Try to run the fused kernel.
+  int fallback = lmha_bwd(bwd_params);
+
+  // If it failed, fallback on separate kernels for K and V.
+  if( fallback ) {
+
+    // Gradient on K.
+
+    // Launch the kernel.
+    set_params(params, values, grad_out, queries, grad_keys);
+    res = lmha<true>(params);
+    if( res ) {
+      return res;
+    }
+
+    // Gradient on V.
+
+    // Launch the kernel.
+    set_params(params, keys, queries, grad_out, grad_values);
+    return lmha<true>(params);
+  }
+
+  // It worked...
+  return 0;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace nvidia
+#endif // #ifdef ENABLE_NVIDIA_OPTIMIZATIONS
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+typedef torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> float_accessor;
+
+#define E_BLOCK_SIZE 8
+
+__global__ void causal_dot_product_kernel(
+    const float_accessor queries,
+    const float_accessor keys,
+    const float_accessor values,
+    float_accessor result,
+    const int N,
+    const int H,
+    const int L,
+    const int E,
+    const int M
+) {
+    int n = blockIdx.y;
+    int h = blockIdx.z;
+
+    int e_start = blockIdx.x * E_BLOCK_SIZE;
+    int m = threadIdx.x % M;
+
+    extern __shared__ float shared_mem[];
+    float* shared_kv = shared_mem;
+
+    for (int e_local = 0; e_local < E_BLOCK_SIZE && e_local + e_start < E; e_local++) {
+      shared_kv[m + e_local * M] = 0;
+    }
+
+    for (int t=0; t<L; t++) {
+      float res = 0;
+      for (int e_local = 0; e_local < E_BLOCK_SIZE && e_local + e_start < E; e_local++) {
+        shared_kv[e_local*M + m] += keys[n][h][t][e_local + e_start] * values[n][h][t][m];
+        res += queries[n][h][t][e_local + e_start] * shared_kv[e_local*M + m];
+      }
+      atomicAdd(
+          &result[n][h][t][m],
+          res
+      );
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void causal_dot_product_(const torch::Tensor queries,
+                         const torch::Tensor keys,
+                         const torch::Tensor values,
+                         torch::Tensor product) {
+    // Make sure that we are using the correct GPU device
+    torch::DeviceGuard _guard(queries.device());
+
+    int N = queries.size(0);
+    int H = queries.size(1);
+    int L = queries.size(2);
+    int E = queries.size(3);
+    int M = values.size(3);
+
+    const int blocks_per_sequence = (E + E_BLOCK_SIZE - 1) / E_BLOCK_SIZE;
+
+    dim3 blockDim(M, 1, 1);
+    dim3 gridDim(blocks_per_sequence, N, H);
+    const int shared_mem_forward = E_BLOCK_SIZE * M * sizeof(float);
+
+    causal_dot_product_kernel<<<gridDim, blockDim, shared_mem_forward>>>(
+      queries.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      keys.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      values.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      product.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      N, H, L, E, M
+    );
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void causal_dot_product(const torch::Tensor queries,
+                        const torch::Tensor keys,
+                        const torch::Tensor values,
+                        torch::Tensor product) {
+#ifdef ENABLE_NVIDIA_OPTIMIZATIONS
+  int fallback = nvidia::lmha_fwd(queries, keys, values, product);
+#else
+  int fallback = 1;
+#endif
+  if( fallback ) {
+    causal_dot_product_(queries, keys, values, product);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define M_BLOCK_SIZE 4
+
+// we need shared memory to store
+// kv
+// Backward direction
+// kv_backwards
+// Shared memory usage
+__global__ void causal_dot_backward_query_key_kernel(
+    const float_accessor queries,
+    const float_accessor keys,
+    const float_accessor values,
+    const float_accessor grad_out,
+    float_accessor grad_queries,
+    float_accessor grad_keys,
+    int N,
+    int H,
+    int L,
+    int E,
+    int M
+) {
+    int n = blockIdx.y;
+    int h = blockIdx.z;
+
+    int m_start = blockIdx.x * M_BLOCK_SIZE;
+    int e = threadIdx.x % E;
+
+    extern __shared__ float shared_mem[];
+    const int shared_kv_size = M_BLOCK_SIZE * E;
+    float* shared_kv = shared_mem;
+    float* shared_kv_bw = shared_mem + shared_kv_size;
+
+    for (int m_local = 0; m_local < M_BLOCK_SIZE && m_local + m_start < M; m_local++) {
+      shared_kv[m_local * E + e] = 0;
+      shared_kv_bw[m_local * E + e] = 0;
+    }
+
+    for (int l=0; l<L; l++) {
+      float res = 0, res_bw = 0;
+      int l_b = L - l - 1;
+      for (int m_local = 0; m_local < M_BLOCK_SIZE && m_local + m_start < M; m_local++) {
+        shared_kv[m_local*E + e] += keys[n][h][l][e] * values[n][h][l][m_start + m_local];
+        shared_kv_bw[m_local*E + e] += queries[n][h][l_b][e] * grad_out[n][h][l_b][m_start + m_local];
+        res += grad_out[n][h][l][m_start + m_local] * shared_kv[m_local*E + e];
+        res_bw += values[n][h][l_b][m_start + m_local] * shared_kv_bw[m_local*E + e];
+      }
+      atomicAdd(
+        &grad_queries[n][h][l][e],
+        res
+      );
+      atomicAdd(
+        &grad_keys[n][h][l_b][e],
+        res_bw
+      );
+    }
+}
+
+
+__global__ void causal_dot_backward_value_kernel(
+    const float_accessor queries,
+    const float_accessor keys,
+    const float_accessor values,
+    const float_accessor grad_out,
+    float_accessor grad_keys,
+    float_accessor grad_values,
+    int N,
+    int H,
+    int L,
+    int E,
+    int M
+) {
+    int n = blockIdx.y;
+    int h = blockIdx.z;
+
+    int e_start = blockIdx.x * E_BLOCK_SIZE;
+    int m = threadIdx.x % M;
+
+    extern __shared__ float shared_mem[];
+    float* shared_kv = shared_mem;
+    for (int e_local = 0; e_local < E_BLOCK_SIZE && e_local + e_start < E; e_local++) {
+      shared_kv[m + e_local * M] = 0;
+    }
+
+    for (int l = 0; l < L; l++) {
+        int l_b = L - l -1;
+        float res = 0;
+        for (int e_local = 0; e_local < E_BLOCK_SIZE && e_local + e_start < E; e_local++) {
+          shared_kv[e_local*M + m] += queries[n][h][l_b][e_start + e_local] * grad_out[n][h][l_b][m];
+          res += keys[n][h][l_b][e_start + e_local] * shared_kv[e_local*M + m];
+        }
+        atomicAdd(
+            &grad_values[n][h][l_b][m],
+            res
+        );
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void causal_dot_backward_(const torch::Tensor queries,
+                          const torch::Tensor keys,
+                          const torch::Tensor values,
+                          const torch::Tensor grad_out,
+                          torch::Tensor grad_queries,
+                          torch::Tensor grad_keys,
+                          torch::Tensor grad_values) {
+
+    // Make sure that we are using the correct GPU device
+    torch::DeviceGuard _guard(queries.device());
+
+    int N = queries.size(0);
+    int H = queries.size(1);
+    int L = queries.size(2);
+    int E = queries.size(3);
+    int M = values.size(3);
+
+    const int blocks_per_sequence = (M + M_BLOCK_SIZE - 1) / M_BLOCK_SIZE;
+
+    dim3 blockDim(E, 1, 1);
+    dim3 gridDim(blocks_per_sequence, N, H);
+    const int shared_mem_qk_backward = 2 * M_BLOCK_SIZE * E * sizeof(float);
+
+    causal_dot_backward_query_key_kernel<<<gridDim, blockDim, shared_mem_qk_backward>>>(
+      queries.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      keys.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      values.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_out.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_queries.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_keys.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      N, H, L, E, M
+    );
+
+    const int blocks_per_sequence_value = (E + E_BLOCK_SIZE - 1) / E_BLOCK_SIZE;
+
+    dim3 blockDimv(M, 1, 1);
+    dim3 gridDimv(blocks_per_sequence_value, N, H);
+    const int shared_mem_v_backward = E_BLOCK_SIZE * M * sizeof(float);
+    causal_dot_backward_value_kernel<<<gridDimv, blockDimv, shared_mem_v_backward>>>(
+      queries.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      keys.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      values.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_out.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_keys.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_values.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      N, H, L, E, M
+    );
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void causal_dot_backward(const torch::Tensor queries,
+                         const torch::Tensor keys,
+                         const torch::Tensor values,
+                         const torch::Tensor grad_out,
+                         torch::Tensor grad_queries,
+                         torch::Tensor grad_keys,
+                         torch::Tensor grad_values) {
+#ifdef ENABLE_NVIDIA_OPTIMIZATIONS
+  int fallback = nvidia::lmha_bwd(queries,
+                                  keys,
+                                  values,
+                                  grad_out,
+                                  grad_queries,
+                                  grad_keys,
+                                  grad_values);
+#else
+  int fallback = 1;
+#endif
+  if( fallback ) {
+    // Make sure that the gradient tensors are 0. This is needed because the
+    // bwd pass might have partially executed and filled in some values in
+    // grad_queries or grad_keys.
+    //
+    // This adds a small overhead every time we have to fall back to the old
+    // kernel for the backward pass.
+    grad_queries.zero_();
+    grad_keys.zero_();
+    causal_dot_backward_(queries, keys, values, grad_out, grad_queries, grad_keys, grad_values);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def(
+        "causal_dot_product",
+        &causal_dot_product,
+        "Compute the weighted sum of values but attending only to previous "
+        "values."
+    );
+    m.def(
+        "causal_dot_backward",
+        &causal_dot_backward,
+        "Compute the gradients for the causal dot product."
+    );
+}

csrc/causal_attention_kv_cuda.cu

@@ -0,0 +1,1483 @@
+//
+// Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/
+// Written by Angelos Katharopoulos <angelos.katharopoulos@idiap.ch>,
+// Apoorv Vyas <avyas@idiap.ch>
+//
+
+//
+// For modifications made inside namespace nvidia (authored by jdemouth):
+//
+// Copyright (c) 2021 NVIDIA CORPORATION. All rights reserved.
+// 
+// Permission is hereby granted, free of charge, to any person obtaining a copy of
+// this software and associated documentation files (the "Software"), to deal in
+// the Software without restriction, including without limitation the rights to
+// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+// the Software, and to permit persons to whom the Software is furnished to do so,
+// subject to the following conditions:
+// 
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+// 
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+// FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+// COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+// IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+// CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+//
+
+#include <torch/extension.h>
+#include <assert.h>
+#include <stdio.h>
+
+#define ENABLE_NVIDIA_OPTIMIZATIONS
+
+#ifdef ENABLE_NVIDIA_OPTIMIZATIONS
+namespace nvidia {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+constexpr int THREADS_PER_WARP = 32;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+constexpr int LOW_OCCUPANCY_THRESHOLD = 40; // TODO: Make it HW specific (like 1/2 SMs).
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ __host__ int div_up(int m, int n) {
+  return (m + n-1) / n;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ __host__ int round_up(int m, int n) {
+  return div_up(m, n) * n;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename T >
+struct Lmha_params {
+
+  // The output buffer. Dimensions [B, H, L, M].
+  T *out;
+
+  // The input Qs. Dimensions [B, H, L, E].
+  const T *q;
+  // The input Ks. Dimensions [B, H, L, E].
+  const T *k;
+  // The input Vs. Dimensions [B, H, L, M].
+  const T *v;
+
+  // The different dimensions.
+  int B, L, H, E, M;
+
+  // The strides for the different tensors.
+  int q_stride_B, q_stride_H, q_stride_L;
+  int k_stride_B, k_stride_H, k_stride_L;
+  int v_stride_B, v_stride_H, v_stride_L;
+  int o_stride_B, o_stride_H, o_stride_L;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, bool GO_BACKWARD, int WARPS, int COLS_PER_THREAD = 4 >
+__global__ __launch_bounds__(WARPS * THREADS_PER_WARP)
+void lmha_low_occupancy_kernel(Lmha_params<float> params) {
+
+  // The number of threads per block.
+  constexpr int THREADS_PER_BLOCK = WARPS * THREADS_PER_WARP;
+  // The number of rows per thread.
+  constexpr int ROWS_PER_THREAD = E / THREADS_PER_WARP;
+  // The number of steps per iteration.
+  constexpr int COLS_PER_ITER = WARPS * COLS_PER_THREAD;
+
+  // Make sure E is a multiple of the warp size.
+  static_assert(E % THREADS_PER_WARP == 0, "");
+
+  // Shared memory to store V/O.
+  __shared__ float smem_v[COLS_PER_ITER], smem_o[COLS_PER_ITER];
+  // Shared memory buffer to performance the reductions.
+  __shared__ float smem_reds[E * WARPS]; 
+
+  // The sequence processed by that block.
+  const int bi = blockIdx.z;
+  // The head processed by that block.
+  const int hi = blockIdx.y;
+  // The hidden cell in the V/output buffers.
+  const int vi = blockIdx.x;
+
+  // The linear index of the thread.
+  const int tidx = threadIdx.x;
+
+  // Decompose the block in warp/lane.
+  const int warp = tidx / THREADS_PER_WARP;
+  const int lane = tidx % THREADS_PER_WARP;
+
+  // The base offset loaded by the thread in Q and K.
+  int offset_q = bi*params.q_stride_B + hi*params.q_stride_H + lane;
+  int offset_k = bi*params.k_stride_B + hi*params.k_stride_H + lane;
+
+  // If we walk backward, account for the extra offset.
+  if( GO_BACKWARD ) {
+    offset_q += (params.L-1)*params.q_stride_L;
+    offset_k += (params.L-1)*params.k_stride_L;
+  }
+
+  // Position the warp at the beginning of the proper timestep.
+  if( GO_BACKWARD ) {
+    offset_q -= warp*COLS_PER_THREAD*params.q_stride_L;
+    offset_k -= warp*COLS_PER_THREAD*params.k_stride_L;
+  } else {
+    offset_q += warp*COLS_PER_THREAD*params.q_stride_L;
+    offset_k += warp*COLS_PER_THREAD*params.k_stride_L;
+  }
+
+  // Determine the base pointers for Q and K.
+  const float *ptr_q = &params.q[offset_q];
+  const float *ptr_k = &params.k[offset_k];
+
+  // Is a given row valid?
+  int valid_qk[ROWS_PER_THREAD];
+  #pragma unroll
+  for( int ii = 0; ii < ROWS_PER_THREAD; ++ii ) {
+    valid_qk[ii] = lane + ii*THREADS_PER_WARP < params.E;
+  }
+
+  // The offset to the position loaded by the thread in V.
+  int offset_v = bi*params.v_stride_B + hi*params.v_stride_H + vi;
+  int offset_o = bi*params.o_stride_B + hi*params.o_stride_H + vi;
+
+  // If we walk backward, account for the extra offset.
+  if( GO_BACKWARD ) {
+    offset_v += (params.L-1)*params.v_stride_L;
+    offset_o += (params.L-1)*params.o_stride_L;
+  }
+
+  // We load/store a strided matrix of COLS_PER_ITER x OUTPUTS_PER_BLOCK.
+  if( GO_BACKWARD ) {
+    offset_v -= tidx*params.v_stride_L;
+    offset_o -= tidx*params.o_stride_L;
+  } else {
+    offset_v += tidx*params.v_stride_L;
+    offset_o += tidx*params.o_stride_L;
+  }
+
+  // Determine the base pointer for V.
+  const float *ptr_v = &params.v[offset_v];
+  // The output pointer. 
+  float *ptr_o = &params.out[offset_o];
+
+  // The running KVs.
+  float running_kv[ROWS_PER_THREAD];
+  #pragma unroll
+  for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+    running_kv[ri] = 0.f;
+  }
+
+  // Iterate over the timesteps. TODO: Use params.loop_count!!!
+  for( int iter = 0; iter < params.L; iter += COLS_PER_ITER ) {
+
+    // Each thread loads a matrix of elements.
+    float q[ROWS_PER_THREAD][COLS_PER_THREAD], k[ROWS_PER_THREAD][COLS_PER_THREAD];
+
+    // Trigger the memory loads for Q and K.
+    #pragma unroll
+    for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+      #pragma unroll
+      for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+
+        // For Q/K, each warp loads from various timesteps. 
+        int ti = iter + warp*COLS_PER_THREAD;
+        if( GO_BACKWARD ) {
+          ti = params.L - 1 - ti;
+        }
+
+        // Is it a valid access?
+        int valid;
+        if( GO_BACKWARD ) {
+          valid = valid_qk[ri] && ti - ci >= 0;
+        } else {
+          valid = valid_qk[ri] && ti + ci < params.L;
+        }
+
+        // The extra offset to add.
+        if( GO_BACKWARD ) {
+          offset_q = ri*THREADS_PER_WARP - ci*params.q_stride_L;
+          offset_k = ri*THREADS_PER_WARP - ci*params.k_stride_L;
+        } else {
+          offset_q = ri*THREADS_PER_WARP + ci*params.q_stride_L;
+          offset_k = ri*THREADS_PER_WARP + ci*params.k_stride_L;
+        }
+
+        // Load Q/K if they are valid.
+        q[ri][ci] = valid ? ptr_q[offset_q] : 0.f;
+        k[ri][ci] = valid ? ptr_k[offset_k] : 0.f;
+      }
+    }
+
+    // For the V tensor, we assign contiguous thread to different loads. So, ti is different.
+    int ti = iter + tidx;
+    if( GO_BACKWARD ) {
+      ti = params.L - 1 - ti;
+    }
+
+    // Is it a valid access?
+    int valid_vo = tidx < COLS_PER_ITER;
+    if( GO_BACKWARD ) {
+      valid_vo &= ti >= 0;
+    } else {
+      valid_vo &= ti < params.L;
+    }
+
+    // Trigger the loads for V. 
+    float ldg_v = valid_vo ? *ptr_v : 0.f;
+
+    // Move the load pointers.
+    if( GO_BACKWARD ) {
+      ptr_q -= COLS_PER_ITER*params.q_stride_L;
+      ptr_k -= COLS_PER_ITER*params.k_stride_L;
+      ptr_v -= COLS_PER_ITER*params.v_stride_L;
+    } else {
+      ptr_q += COLS_PER_ITER*params.q_stride_L;
+      ptr_k += COLS_PER_ITER*params.k_stride_L;
+      ptr_v += COLS_PER_ITER*params.v_stride_L;
+    }
+
+    // Store to shared memory.
+    if( tidx < COLS_PER_ITER ) {
+      smem_v[tidx] = ldg_v;
+    }
+
+    // Make sure V is in shared memory.
+    __syncthreads();
+
+    // Read V from shared memory.
+    float v[COLS_PER_THREAD];
+    #pragma unroll
+    for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+      v[ci] = smem_v[warp*COLS_PER_THREAD + ci];
+    }
+
+    // Each thread computes local K*V products.
+    float kv[ROWS_PER_THREAD][COLS_PER_THREAD];
+    #pragma unroll
+    for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+      #pragma unroll
+      for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+        kv[ri][ci] = 0.f;
+      }
+    }
+
+    // Update the K*V^T product.
+    #pragma unroll
+    for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+      #pragma unroll
+      for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+        kv[ri][ci] += k[ri][ci] * v[ci];
+      }
+    }
+
+    // We must perform the prefix sums within the thread-block. Start with the thread.
+    #pragma unroll
+    for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+      #pragma unroll
+      for( int ci = 1; ci < COLS_PER_THREAD; ++ci ) {
+        kv[ri][ci] += kv[ri][ci-1];
+      }
+    }
+
+    // Store the partial sums to shared memory. Unless we have no inter-warp reduction to perform.
+    #pragma unroll
+    for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+      smem_reds[warp*E + ri*THREADS_PER_WARP + lane] = kv[ri][COLS_PER_THREAD-1];
+    }
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Each thread deals with one or more column(s) of the matrix.
+    constexpr int SUMS_PER_THREAD = (E + THREADS_PER_BLOCK-1) / THREADS_PER_BLOCK;
+    #pragma unroll
+    for( int ii = 0, idx = tidx; ii < SUMS_PER_THREAD; ++ii, idx += THREADS_PER_BLOCK ) {
+      if( idx < E ) {
+        float sum = smem_reds[idx];
+        #pragma unroll
+        for( int jj = 1; jj < WARPS; ++jj ) {
+          smem_reds[idx + jj*E] = sum += smem_reds[idx + jj*E];
+        }
+      }
+    }
+
+    // Make sure the reductions are stored in shared memory.
+    __syncthreads();
+
+    // Each thread updates his partial products.
+    #pragma unroll
+    for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+      float sum = running_kv[ri];
+      if( warp > 0 ) {
+        sum += smem_reds[(warp-1)*E + lane + ri*THREADS_PER_WARP];
+      }
+      #pragma unroll
+      for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+        kv[ri][ci] += sum;
+      }
+    }
+
+    // Compute the partial output values for that thread.
+    float sum[COLS_PER_THREAD];
+    #pragma unroll
+    for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+      sum[ci] = q[0][ci] * kv[0][ci];
+      #pragma unroll
+      for( int ri = 1; ri < ROWS_PER_THREAD; ++ri ) {
+        sum[ci] += q[ri][ci] * kv[ri][ci];
+      }
+    }
+
+    // Run the parallel reductions inside the warp.
+    #pragma unroll
+    for( int mask = THREADS_PER_WARP / 2; mask >= 1; mask /= 2 ) {
+      #pragma unroll
+      for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+        sum[ci] += __shfl_xor_sync(uint32_t(-1), sum[ci], mask);
+      }
+    }
+
+    // Store the final output to shared memory.
+    if( lane == 0 ) {
+      #pragma unroll
+      for( int ci = 0; ci < COLS_PER_THREAD; ++ci ) {
+        smem_o[warp*COLS_PER_THREAD + ci] = sum[ci];
+      }
+    }
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Store the output.
+    if( valid_vo ) {
+      *ptr_o = smem_o[tidx];
+    }
+
+    // Each thread updates his running kv.
+    #pragma unroll
+    for( int ri = 0; ri < ROWS_PER_THREAD; ++ri ) {
+      running_kv[ri] += smem_reds[(WARPS-1)*E + lane + ri*THREADS_PER_WARP];
+    }
+
+    // Move to next location.
+    if( GO_BACKWARD ) {
+      ptr_o -= COLS_PER_ITER*params.o_stride_L;
+    } else {
+      ptr_o += COLS_PER_ITER*params.o_stride_L;
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, bool GO_BACKWARD, int WARPS >
+int lmha_low_occupancy_(const Lmha_params<float> &params) {
+
+  // Make sure we are not going to launch an invalid grid.
+  if( params.H > 65535 || params.B > 65535 ) {
+    return 1;
+  }
+
+  // Prepare the grid and trigger the CUDA kernel.
+  dim3 grid;
+  grid.x = params.M;
+  grid.y = params.H;
+  grid.z = params.B;
+  lmha_low_occupancy_kernel<E, GO_BACKWARD, WARPS><<<grid, WARPS*THREADS_PER_WARP>>>(params);
+  return 0;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, bool GO_BACKWARD >
+int lmha_low_occupancy_(const Lmha_params<float> &params, int blocks) {
+         if( params.M * blocks >= 8*LOW_OCCUPANCY_THRESHOLD ) {
+    return lmha_low_occupancy_<E, GO_BACKWARD,  4>(params);
+  } else if( params.M * blocks >= 4*LOW_OCCUPANCY_THRESHOLD ) {
+    return lmha_low_occupancy_<E, GO_BACKWARD,  8>(params);
+  } else {
+    return lmha_low_occupancy_<E, GO_BACKWARD, 16>(params);
+  }
+  return 1;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, typename Params >
+static inline __device__ __host__ int smem_buffer_elts_(const Params &params) {
+  int M = round_up(params.M, 4);
+  return 2*E + 2*M;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, int THREADS_PER_HEAD, bool GO_BACKWARD >
+__global__ 
+void lmha_kernel(Lmha_params<float> params) {
+
+  // Make sure E is a multiple of 4.
+  static_assert(E % 4 == 0, "");
+
+  // The amount of shared memory per buffer (2 buffers for double-buffering).
+  const int smem_buffer_elts = smem_buffer_elts_<E>(params);
+  // The M dimension for shared memory.
+  const int M = round_up(params.M, 4);
+
+  // Shared memory to store Q, K and V. Size is 2*smem_buffer_elts.
+  extern __shared__ float smem_[];
+
+  // The various shared memory buffers.
+  float *smem_q = &smem_[0*E];
+  float *smem_k = &smem_[1*E];
+  float *smem_v = &smem_[2*E];
+  float *smem_o = &smem_[2*E + M];
+
+  // The index of the shared memory buffer (for double-buffering).
+  int smem_curr = 0;
+
+  // The sequence processed by that block.
+  const int bi = blockIdx.y;
+  // The head processed by that block.
+  const int hi = blockIdx.x;
+
+  // The linear index of the thread.
+  const int tidx = threadIdx.x;
+
+  // The offset to the position loaded by the thread in Q.
+  int offset_q = bi*params.q_stride_B + hi*params.q_stride_H + tidx;
+  // The offset to the position loaded by the thread in K.
+  int offset_k = bi*params.k_stride_B + hi*params.k_stride_H + tidx;
+
+  // If we walk backward, account for the extra offset.
+  if( GO_BACKWARD ) {
+    offset_q += (params.L-1)*params.q_stride_L;
+    offset_k += (params.L-1)*params.k_stride_L;
+  }
+
+  // Determine the base pointers for Q and K.
+  const float *ptr_q = &params.q[offset_q];
+  const float *ptr_k = &params.k[offset_k];
+
+  // The offset to the position loaded by the thread in V and O.
+  int offset_v = bi*params.v_stride_B + hi*params.v_stride_H + tidx;
+  int offset_o = bi*params.o_stride_B + hi*params.o_stride_H + tidx;
+
+  // If we walk backward, account for the extra offset.
+  if( GO_BACKWARD ) {
+    offset_v += (params.L-1)*params.v_stride_L;
+    offset_o += (params.L-1)*params.o_stride_L;
+  }
+
+  // Determine the base pointers for V.
+  const float *ptr_v = &params.v[offset_v];
+
+  // Is it an active Q/K thread?
+  const int active_qk = tidx < params.E;
+
+  // Trigger the memory loads for Q and K.
+  float ldg_q = 0.f, ldg_k = 0.f;
+  if( active_qk ) {
+    ldg_q = *ptr_q;
+    ldg_k = *ptr_k;
+  }
+
+  // Is it an active V thread?
+  const int active_v = tidx < params.M;
+
+  // Trigger the memory loads for V. 
+  float ldg_v = 0.f;
+  if( active_v ) {
+    ldg_v = *ptr_v;
+  }
+
+  // Move the load pointers.
+  if( GO_BACKWARD ) {
+    ptr_q -= params.q_stride_L;
+    ptr_k -= params.k_stride_L;
+    ptr_v -= params.v_stride_L;
+  } else {
+    ptr_q += params.q_stride_L;
+    ptr_k += params.k_stride_L;
+    ptr_v += params.v_stride_L;
+  }
+
+  // The number of FLOAT4s per head.
+  constexpr int FLOAT4s_PER_HEAD = E / 4;
+  // The number of FLOAT4s per thread.
+  constexpr int FLOAT4s_PER_THREAD = FLOAT4s_PER_HEAD / THREADS_PER_HEAD;
+
+  // The storage for the K*V^T values.
+  float4 kv[FLOAT4s_PER_THREAD]; 
+  #pragma unroll
+  for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+    kv[ii] = make_float4(0.f, 0.f, 0.f, 0.f);
+  }
+
+  // The output pointer.
+  float *out_ptr = &params.out[offset_o];
+
+  // Store to shared memory Q and K.
+  if( tidx < E ) { 
+    smem_q[smem_curr*smem_buffer_elts + tidx] = ldg_q; 
+    smem_k[smem_curr*smem_buffer_elts + tidx] = ldg_k; 
+  }
+
+  // Store to shared memory V. All threads store valid values.
+  if( tidx < M ) {
+    smem_v[smem_curr*smem_buffer_elts + tidx] = ldg_v;
+  }
+
+  // The position of the thread in the V dimension.
+  int vo = tidx / THREADS_PER_HEAD;
+  int vi = tidx % THREADS_PER_HEAD;
+
+  // Iterate over the timesteps.
+  for( int ti = 0; ti < params.L; ++ti ) {
+
+    // Is it the last iteration?
+    int is_last = ti == params.L - 1;
+
+    // Trigger the next loads for Q and K.
+    if( !is_last && active_qk ) {
+      ldg_q = *ptr_q;
+      ldg_k = *ptr_k;
+    }
+
+    // Trigger the next loads for V.
+    if( !is_last && active_v ) {
+      ldg_v = *ptr_v;
+    }
+
+    // Move the load pointers.
+    if( GO_BACKWARD ) {
+      ptr_q -= params.q_stride_L;
+      ptr_k -= params.k_stride_L;
+      ptr_v -= params.v_stride_L;
+    } else {
+      ptr_q += params.q_stride_L;
+      ptr_k += params.k_stride_L;
+      ptr_v += params.v_stride_L;
+    }
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Each thread loads 4 values from K.
+    float4 k[FLOAT4s_PER_THREAD];
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      int ki = tidx % THREADS_PER_HEAD * 4 + ii * THREADS_PER_HEAD * 4;
+      k[ii] = *reinterpret_cast<const float4*>(&smem_k[smem_curr*smem_buffer_elts + ki]);
+    }
+
+    // Each thread loads a single V value.
+    float v = 0.f;
+    if( vo < params.M ) {
+      v = *reinterpret_cast<const float *>(&smem_v[smem_curr*smem_buffer_elts + vo]);
+    }
+
+    // Update the K*V^T product.
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      kv[ii].x += k[ii].x * v;
+      kv[ii].y += k[ii].y * v;
+      kv[ii].z += k[ii].z * v;
+      kv[ii].w += k[ii].w * v;
+    }
+
+    // Load the Q values from shared memory.
+    float4 q[FLOAT4s_PER_THREAD]; 
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      int qi = tidx % THREADS_PER_HEAD * 4 + ii * THREADS_PER_HEAD * 4;
+      q[ii] = *reinterpret_cast<const float4*>(&smem_q[smem_curr*smem_buffer_elts + qi]);
+    }
+
+    // Compute the partial output value for that thread.
+    float sum = 0.f;
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      sum += q[ii].x * kv[ii].x;
+      sum += q[ii].y * kv[ii].y;
+      sum += q[ii].z * kv[ii].z;
+      sum += q[ii].w * kv[ii].w;
+    }
+
+    // Finalize the computation of the sum (if we have more than 1 thread per head).
+    if( THREADS_PER_HEAD > 1 ) {
+
+      // Finalize the sum for each head.
+      #pragma unroll
+      for( int mask = THREADS_PER_HEAD / 2; mask >= 1; mask /= 2 ) {
+        sum += __shfl_xor_sync(uint32_t(-1), sum, mask);
+      }
+
+      // Store to shared memory.
+      if( vo < M && vi == 0 ) {
+        smem_o[smem_curr*smem_buffer_elts + vo] = sum;
+      }
+
+      // Make sure the data is in shared memory.
+      __syncthreads();
+
+      // Active threads read the data to store.
+      if( active_v ) {
+        sum = smem_o[smem_curr*smem_buffer_elts + tidx];
+      }
+
+    } // THREADS_PER_HEAD > 1.
+
+    // Store the output. All the threads are active.
+    if( active_v ) {
+      *out_ptr = sum;
+    }
+
+    // Move to next location.
+    if( GO_BACKWARD ) {
+      out_ptr -= params.o_stride_L;
+    } else {
+      out_ptr += params.o_stride_L;
+    }
+
+    // Move the shared memory buffer.
+    smem_curr = (smem_curr + 1) % 2;
+
+    // Store to shared memory for Q and K.
+    if( !is_last && tidx < E ) {
+      smem_q[smem_curr*smem_buffer_elts + tidx] = ldg_q;
+      smem_k[smem_curr*smem_buffer_elts + tidx] = ldg_k;
+    }
+
+    // Store to shared memory for V.
+    if( !is_last && tidx < M ) {
+      smem_v[smem_curr*smem_buffer_elts + tidx] = ldg_v;
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int E, int THREADS_PER_HEAD, bool GO_BACKWARD >
+int lmha_(const Lmha_params<float> &params) {
+  // The M dimension rounded up to 4.
+  int M = round_up(params.M, 4);
+
+  // The number of threads in the block.
+  int block = round_up(max(E, M*THREADS_PER_HEAD), 32);
+  if( block > 512 || params.B > 65535 ) {
+    return 1;
+  }
+
+  // Prepare the kernel.
+  dim3 grid(params.H, params.B);
+  size_t smem = smem_buffer_elts_<E>(params)*2*sizeof(float);
+  lmha_kernel<E, THREADS_PER_HEAD, GO_BACKWARD><<<grid, block, smem>>>(params);
+  return 0;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< bool GO_BACKWARD >
+int lmha(const Lmha_params<float> &params) {
+  int blocks = params.B * params.H;
+  int res = 1;
+  if( blocks < LOW_OCCUPANCY_THRESHOLD ) { 
+           if( params.E <=  32 ) {
+      res = lmha_low_occupancy_< 32, GO_BACKWARD>(params, blocks);
+    } else if( params.E <=  64 ) {
+      res = lmha_low_occupancy_< 64, GO_BACKWARD>(params, blocks);
+    } else if( params.E <= 128 ) {
+      res = lmha_low_occupancy_<128, GO_BACKWARD>(params, blocks);
+    } else if( params.E <= 256 ) {
+      res = lmha_low_occupancy_<256, GO_BACKWARD>(params, blocks);
+    }
+  } else {
+           if( params.E <=  32 ) {
+      res = lmha_< 32, 1, GO_BACKWARD>(params);
+    } else if( params.E <=  48 ) {
+      res = lmha_< 48, 1, GO_BACKWARD>(params);
+    } else if( params.E <=  64 ) {
+      res = lmha_< 64, 1, GO_BACKWARD>(params);
+    } else if( params.E <= 128 ) {
+      res = lmha_<128, 2, GO_BACKWARD>(params);
+    } else if( params.E <= 256 ) {
+      res = lmha_<256, 4, GO_BACKWARD>(params);
+    }
+  }
+  return res;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename T >
+inline void set_params(Lmha_params<T> &params,
+                       const torch::Tensor q,
+                       const torch::Tensor k,
+                       const torch::Tensor v,
+                       torch::Tensor       o) {
+
+  // Define the pointers.
+  params.out = o.data_ptr<T>();
+  params.q   = q.data_ptr<T>();
+  params.k   = k.data_ptr<T>();
+  params.v   = v.data_ptr<T>();
+
+  // Define the strides.
+  params.q_stride_B = (int) q.stride(0);
+  params.q_stride_H = (int) q.stride(1);
+  params.q_stride_L = (int) q.stride(2);
+  params.k_stride_B = (int) k.stride(0);
+  params.k_stride_H = (int) k.stride(1);
+  params.k_stride_L = (int) k.stride(2);
+  params.v_stride_B = (int) v.stride(0);
+  params.v_stride_H = (int) v.stride(1);
+  params.v_stride_L = (int) v.stride(2);
+  params.o_stride_B = (int) o.stride(0);
+  params.o_stride_H = (int) o.stride(1);
+  params.o_stride_L = (int) o.stride(2);
+
+  // Extract the dimensions.
+  int N = q.size(0);
+  int H = q.size(1);
+  int L = q.size(2);
+  int E = q.size(3);
+  int M = v.size(3);
+
+  params.B = N;
+  params.L = L;
+  params.H  = H;
+  params.E = E;
+  params.M = M;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+int lmha_fwd(const torch::Tensor queries,
+             const torch::Tensor keys,
+             const torch::Tensor values,
+             torch::Tensor product) {
+
+  // Make sure that we are using the correct GPU device
+  torch::DeviceGuard _guard(queries.device());
+
+  // Make sure the inner-most dimension of the tensors is packed.
+  assert(queries.stride(3) == 1);
+  assert(keys   .stride(3) == 1);
+  assert(values .stride(3) == 1);
+  assert(product.stride(3) == 1);
+
+  // Extract the dimensions.
+  int N = queries.size(0);
+  int H = queries.size(1);
+  int L = queries.size(2);
+  int E = queries.size(3);
+  int M = values.size (3);
+
+  // The structure of params.
+  Lmha_params<float> params;
+  set_params(params, queries, keys, values, product);
+
+  // Launch the kernel.
+  return lmha<false>(params);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename T >
+struct Lmha_bwd_params {
+
+  // The output buffer for K. Dimensions [B, H, L, D].
+  T *out_k;
+  // The output buffer for V. Dimensions [B, H, L, D].
+  T *out_v;
+
+  // The input Qs. Dimensions [B, H, L, D].
+  const T *q;
+  // The input Ks. Dimensions [B, H, L, D].
+  const T *k;
+  // The input Vs. Dimensions [B, H, L, D].
+  const T *v;
+  // The input Gs. Dimensions [B, H, L, D].
+  const T *g;
+
+  // The dimensions.
+  int B, L, H, M, E;
+
+  // The strides for the input tensors.
+  int q_stride_B, q_stride_L, q_stride_H;
+  int k_stride_B, k_stride_L, k_stride_H;
+  int v_stride_B, v_stride_L, v_stride_H;
+  int g_stride_B, g_stride_L, g_stride_H;
+
+  // The strides for the outputs.
+  int out_k_stride_B, out_k_stride_L, out_k_stride_H;
+  int out_v_stride_B, out_v_stride_L, out_v_stride_H;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int D, int THREADS_PER_HEAD >
+__global__ __launch_bounds__(D*THREADS_PER_HEAD*2)
+void lmha_bwd_kernel(Lmha_bwd_params<float> params) {
+
+  // Make sure D is a multiple of 4.
+  static_assert(D % 4 == 0, "");
+
+  // The shared memory buffers.
+  __shared__ struct Smem { float qg[2*D], kv[2*D], out_kv[2*D]; } smem_[2];
+
+  // The index of the shared memory buffer (for double-buffering).
+  int smem_curr = 0;
+
+  // The sequence processed by that block.
+  const int bi = blockIdx.y;
+  // The head processed by that block.
+  const int hi = blockIdx.x;
+
+  // The linear index of the thread.
+  const int tidx = threadIdx.x;
+
+  // Split the threads into two slices.
+  int so = tidx / (D*THREADS_PER_HEAD);
+  int si = tidx % (D*THREADS_PER_HEAD);
+
+  // The strides for B/L/H for the Q/G tensors.
+  int qg_stride_B, qg_stride_L, qg_stride_H;
+  if( so == 0 ) {
+    qg_stride_B = params.q_stride_B;
+    qg_stride_L = params.q_stride_L;
+    qg_stride_H = params.q_stride_H;
+  } else {
+    qg_stride_B = params.g_stride_B;
+    qg_stride_L = params.g_stride_L;
+    qg_stride_H = params.g_stride_H;
+  }
+
+  // The strides for B/L/H for the K/V tensors.
+  int kv_stride_B, kv_stride_L, kv_stride_H;
+  if( so == 0 ) {
+    kv_stride_B = params.k_stride_B;
+    kv_stride_L = params.k_stride_L;
+    kv_stride_H = params.k_stride_H;
+  } else {
+    kv_stride_B = params.v_stride_B;
+    kv_stride_L = params.v_stride_L;
+    kv_stride_H = params.v_stride_H;
+  }
+
+  // The hidden size.
+  int hidden_size_per_head = 0;
+  if( so == 0 ) {
+    hidden_size_per_head = params.E;
+  } else {
+    hidden_size_per_head = params.M;
+  }
+
+  // Where to start reading from.
+  int offset_qg = bi*qg_stride_B + hi*qg_stride_H + si;
+  int offset_kv = bi*kv_stride_B + hi*kv_stride_H + si;
+
+  // We walk backward, account for the extra offset.
+  offset_qg += (params.L-1)*qg_stride_L;
+  offset_kv += (params.L-1)*kv_stride_L;
+
+  // Determine the base pointers for Q, K, V and G.
+  const float *ptr_qg = &(so == 0 ? params.q : params.g)[offset_qg];
+  const float *ptr_kv = &(so == 0 ? params.k : params.v)[offset_kv]; 
+
+  // Is it an active thread?
+  const int active = si < hidden_size_per_head;
+
+  // Trigger the memory loads for Q, K, V and G.
+  float ldg_qg = 0.f, ldg_kv = 0.f;
+  if( active ) {
+    ldg_qg = *ptr_qg;
+    ldg_kv = *ptr_kv;
+  }
+
+  // Move the load pointers (backward).
+  ptr_qg -= qg_stride_L;
+  ptr_kv -= kv_stride_L;
+
+  // The number of FLOAT4s per head.
+  constexpr int FLOAT4s_PER_HEAD = D / 4;
+  // The number of FLOAT4s per thread.
+  constexpr int FLOAT4s_PER_THREAD = FLOAT4s_PER_HEAD / THREADS_PER_HEAD;
+
+  // The storage for the G*Q^T or Q^T*G values.
+  float4 gq[FLOAT4s_PER_THREAD]; 
+  #pragma unroll
+  for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+    gq[ii] = make_float4(0.f, 0.f, 0.f, 0.f);
+  }
+
+  // The strides for B/L/H for the K/V tensors.
+  int out_kv_stride_B, out_kv_stride_L, out_kv_stride_H;
+  if( so == 0 ) {
+    out_kv_stride_B = params.out_k_stride_B;
+    out_kv_stride_L = params.out_k_stride_L;
+    out_kv_stride_H = params.out_k_stride_H;
+  } else {
+    out_kv_stride_B = params.out_v_stride_B;
+    out_kv_stride_L = params.out_v_stride_L;
+    out_kv_stride_H = params.out_v_stride_H;
+  }
+
+  // Where to start reading from.
+  int offset_out_kv = bi*out_kv_stride_B + hi*out_kv_stride_H + si;
+
+  // We walk backward, account for the extra offset.
+  offset_out_kv += (params.L-1)*out_kv_stride_L;
+
+  // The output pointer.
+  float *ptr_out_kv = &(so == 0 ? params.out_k : params.out_v)[offset_out_kv];
+
+  // Store to shared memory.
+  if( si < D ) { 
+    smem_[smem_curr].qg[so*D + si] = ldg_qg; 
+    smem_[smem_curr].kv[so*D + si] = ldg_kv; 
+  }
+
+  // The position of the thread in the output dimension.
+  int oo = si / THREADS_PER_HEAD % D;
+  int oi = si % THREADS_PER_HEAD * 4;
+
+  // Iterate over the timesteps.
+  for( int ti = 0; ti < params.L; ++ti ) {
+
+    // Is it the last iteration?
+    int is_last = ti == params.L - 1;
+
+    // Trigger the next loads.
+    if( !is_last && active ) {
+      ldg_qg = *ptr_qg;
+      ldg_kv = *ptr_kv;
+    }
+
+    // Move the load pointers.
+    ptr_qg -= qg_stride_L;
+    ptr_kv -= kv_stride_L;
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Each thread loads 4 values from G or Q.
+    float4 g[FLOAT4s_PER_THREAD];
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      float *smem_ptr = &smem_[smem_curr].qg[(so^1)*D + oi];
+      g[ii] = *reinterpret_cast<const float4*>(&smem_ptr[ii*THREADS_PER_HEAD*4]);
+    }
+
+    // Each thread loads a single from Q or G value.
+    float q = smem_[smem_curr].qg[so*D + oo];
+
+    // Update the G*Q^T or Q*G^T product.
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      gq[ii].x += g[ii].x * q;
+      gq[ii].y += g[ii].y * q;
+      gq[ii].z += g[ii].z * q;
+      gq[ii].w += g[ii].w * q;
+    }
+
+    // Load the V or K values from shared memory.
+    float4 v[FLOAT4s_PER_THREAD]; 
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      float *smem_ptr = &smem_[smem_curr].kv[(so^1)*D + oi];
+      v[ii] = *reinterpret_cast<const float4*>(&smem_ptr[ii*THREADS_PER_HEAD*4]);
+    }
+
+    // Compute the partial output value for that thread.
+    float sum = 0.f;
+    #pragma unroll
+    for( int ii = 0; ii < FLOAT4s_PER_THREAD; ++ii ) {
+      sum += v[ii].x * gq[ii].x;
+      sum += v[ii].y * gq[ii].y;
+      sum += v[ii].z * gq[ii].z;
+      sum += v[ii].w * gq[ii].w;
+    }
+
+    // Finalize the computation of the sum (if we have more than 1 thread per head).
+    if( THREADS_PER_HEAD > 1 ) {
+
+      // Finalize the sum for each head.
+      #pragma unroll
+      for( int mask = THREADS_PER_HEAD / 2; mask >= 1; mask /= 2 ) {
+        sum += __shfl_xor_sync(uint32_t(-1), sum, mask);
+      }
+
+      // Store to shared memory.
+      if( oi == 0 ) {
+        smem_[smem_curr].out_kv[so*D + oo] = sum;
+      }
+
+      // Make sure the data is in shared memory.
+      __syncthreads();
+
+      // Active threads read the data to store.
+      if( si < hidden_size_per_head ) {
+        sum = smem_[smem_curr].out_kv[so*D + si];
+      }
+
+    } // THREADS_PER_HEAD > 1.
+
+    // Store the output. All the threads are active.
+    if( si < hidden_size_per_head ) {
+      *ptr_out_kv = sum;
+    }
+
+    // Move to next location.
+    ptr_out_kv -= out_kv_stride_L;
+
+    // Move the shared memory buffer.
+    smem_curr = (smem_curr + 1) % 2;
+
+    // Store to shared memory for Q and K.
+    if( !is_last && si < D ) {
+      smem_[smem_curr].qg[so*D + si] = ldg_qg; 
+      smem_[smem_curr].kv[so*D + si] = ldg_kv; 
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int D, int THREADS_PER_HEAD >
+int lmha_bwd_(const Lmha_bwd_params<float> &params) {
+  int block = D*THREADS_PER_HEAD*2;
+  if( block >= 1024 || params.B > 65535 ) {
+    return 1;
+  }
+  dim3 grid(params.H, params.B);
+  lmha_bwd_kernel<D, THREADS_PER_HEAD><<<grid, block>>>(params);
+  return 0;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+int lmha_bwd(const Lmha_bwd_params<float> &params) {
+  int blocks = params.B * params.H;
+  if( blocks < LOW_OCCUPANCY_THRESHOLD ) { 
+    return 1;
+  }
+
+  int hidden_size_per_head = max(params.E, params.M);
+  int res = 1;
+  if( hidden_size_per_head <= 32 ) {
+    res = lmha_bwd_< 32, 1>(params);
+  } else if( hidden_size_per_head <= 64 ) {
+    res = lmha_bwd_< 64, 1>(params);
+  } else if( hidden_size_per_head <= 128 ) {
+    res = lmha_bwd_<128, 2>(params);
+  } else if( hidden_size_per_head <= 256 ) {
+    res = lmha_bwd_<256, 4>(params);
+  }
+  return res;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+int lmha_bwd(const torch::Tensor queries,
+             const torch::Tensor keys,
+             const torch::Tensor values,
+             const torch::Tensor grad_out,
+             torch::Tensor grad_queries,
+             torch::Tensor grad_keys,
+             torch::Tensor grad_values) {
+
+  // Make sure that we are using the correct GPU device
+  torch::DeviceGuard _guard(queries.device());
+
+  // Make sure the inner-most dimension of the tensors is packed.
+  assert(queries     .stride(3) == 1);
+  assert(keys        .stride(3) == 1);
+  assert(values      .stride(3) == 1);
+  assert(grad_out    .stride(3) == 1);
+  assert(grad_queries.stride(3) == 1);
+  assert(grad_keys   .stride(3) == 1);
+  assert(grad_values .stride(3) == 1);
+
+  // Extract the dimensions.
+  int N = queries.size(0);
+  int H = queries.size(1);
+  int L = queries.size(2);
+  int E = queries.size(3);
+  int M = values.size (3);
+
+  // Gradient on Q.
+
+  // The structure of params.
+  Lmha_params<float> params;
+  set_params(params, grad_out, values, keys, grad_queries);
+
+  // Launch the kernel.
+  int res = lmha<false>(params);
+  if( res ) {
+    return res;
+  }
+
+  // Gradient on K and V together.
+
+  Lmha_bwd_params<float> bwd_params;
+  bwd_params.out_k = grad_keys.data_ptr<float>();
+  bwd_params.out_v = grad_values.data_ptr<float>();
+  bwd_params.q = queries.data_ptr<float>();
+  bwd_params.k = keys.data_ptr<float>();
+  bwd_params.v = values.data_ptr<float>();
+  bwd_params.g = grad_out.data_ptr<float>();
+
+  bwd_params.B = N;
+  bwd_params.L = L;
+  bwd_params.H = H;
+  bwd_params.E = E;
+  bwd_params.M = M;
+
+  bwd_params.q_stride_B = queries.stride(0);
+  bwd_params.q_stride_H = queries.stride(1);
+  bwd_params.q_stride_L = queries.stride(2);
+  bwd_params.k_stride_B = keys.stride(0);
+  bwd_params.k_stride_H = keys.stride(1);
+  bwd_params.k_stride_L = keys.stride(2);
+  bwd_params.v_stride_B = values.stride(0);
+  bwd_params.v_stride_H = values.stride(1);
+  bwd_params.v_stride_L = values.stride(2);
+  bwd_params.g_stride_B = grad_out.stride(0);
+  bwd_params.g_stride_H = grad_out.stride(1);
+  bwd_params.g_stride_L = grad_out.stride(2);
+
+  bwd_params.out_k_stride_B = grad_keys.stride(0);
+  bwd_params.out_k_stride_H = grad_keys.stride(1);
+  bwd_params.out_k_stride_L = grad_keys.stride(2);
+  bwd_params.out_v_stride_B = grad_values.stride(0);
+  bwd_params.out_v_stride_H = grad_values.stride(1);
+  bwd_params.out_v_stride_L = grad_values.stride(2);
+
+  // Try to run the fused kernel.
+  int fallback = lmha_bwd(bwd_params);
+
+  // If it failed, fallback on separate kernels for K and V.
+  if( fallback ) {
+
+    // Gradient on K.
+
+    // Launch the kernel.
+    set_params(params, values, grad_out, queries, grad_keys);
+    res = lmha<true>(params);
+    if( res ) {
+      return res;
+    }
+
+    // Gradient on V.
+
+    // Launch the kernel.
+    set_params(params, keys, queries, grad_out, grad_values);
+    return lmha<true>(params);
+  }
+
+  // It worked...
+  return 0;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace nvidia
+#endif // #ifdef ENABLE_NVIDIA_OPTIMIZATIONS
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+typedef torch::PackedTensorAccessor32<float, 4, torch::RestrictPtrTraits> float_accessor;
+
+#define E_BLOCK_SIZE 8
+
+__global__ void causal_dot_product_kernel(
+    const float_accessor queries,
+    const float_accessor keys,
+    const float_accessor values,
+    float_accessor result,
+    const int N,
+    const int H,
+    const int L,
+    const int E,
+    const int M
+) {
+    int n = blockIdx.y;
+    int h = blockIdx.z;
+
+    int e_start = blockIdx.x * E_BLOCK_SIZE;
+    int m = threadIdx.x % M;
+
+    extern __shared__ float shared_mem[];
+    float* shared_kv = shared_mem;
+
+    for (int e_local = 0; e_local < E_BLOCK_SIZE && e_local + e_start < E; e_local++) {
+      shared_kv[m + e_local * M] = 0;
+    }
+
+    for (int t=0; t<L; t++) {
+      float res = 0;
+      for (int e_local = 0; e_local < E_BLOCK_SIZE && e_local + e_start < E; e_local++) {
+        shared_kv[e_local*M + m] += keys[n][h][t][e_local + e_start] * values[n][h][t][m];
+        res += queries[n][h][t][e_local + e_start] * shared_kv[e_local*M + m];
+      }
+      atomicAdd(
+          &result[n][h][t][m],
+          res
+      );
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void causal_dot_product_(const torch::Tensor queries,
+                         const torch::Tensor keys,
+                         const torch::Tensor values,
+                         torch::Tensor product) {
+    // Make sure that we are using the correct GPU device
+    torch::DeviceGuard _guard(queries.device());
+
+    int N = queries.size(0);
+    int H = queries.size(1);
+    int L = queries.size(2);
+    int E = queries.size(3);
+    int M = values.size(3);
+
+    const int blocks_per_sequence = (E + E_BLOCK_SIZE - 1) / E_BLOCK_SIZE;
+
+    dim3 blockDim(M, 1, 1);
+    dim3 gridDim(blocks_per_sequence, N, H);
+    const int shared_mem_forward = E_BLOCK_SIZE * M * sizeof(float);
+
+    causal_dot_product_kernel<<<gridDim, blockDim, shared_mem_forward>>>(
+      queries.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      keys.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      values.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      product.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      N, H, L, E, M
+    );
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void causal_dot_product(const torch::Tensor queries,
+                        const torch::Tensor keys,
+                        const torch::Tensor values,
+                        torch::Tensor product) {
+#ifdef ENABLE_NVIDIA_OPTIMIZATIONS
+  int fallback = nvidia::lmha_fwd(queries, keys, values, product);
+#else
+  int fallback = 1;
+#endif
+  if( fallback ) {
+    causal_dot_product_(queries, keys, values, product);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define M_BLOCK_SIZE 4
+
+// we need shared memory to store
+// kv
+// Backward direction
+// kv_backwards
+// Shared memory usage
+__global__ void causal_dot_backward_query_key_kernel(
+    const float_accessor queries,
+    const float_accessor keys,
+    const float_accessor values,
+    const float_accessor grad_out,
+    float_accessor grad_queries,
+    float_accessor grad_keys,
+    int N,
+    int H,
+    int L,
+    int E,
+    int M
+) {
+    int n = blockIdx.y;
+    int h = blockIdx.z;
+
+    int m_start = blockIdx.x * M_BLOCK_SIZE;
+    int e = threadIdx.x % E;
+
+    extern __shared__ float shared_mem[];
+    const int shared_kv_size = M_BLOCK_SIZE * E;
+    float* shared_kv = shared_mem;
+    float* shared_kv_bw = shared_mem + shared_kv_size;
+
+    for (int m_local = 0; m_local < M_BLOCK_SIZE && m_local + m_start < M; m_local++) {
+      shared_kv[m_local * E + e] = 0;
+      shared_kv_bw[m_local * E + e] = 0;
+    }
+
+    for (int l=0; l<L; l++) {
+      float res = 0, res_bw = 0;
+      int l_b = L - l - 1;
+      for (int m_local = 0; m_local < M_BLOCK_SIZE && m_local + m_start < M; m_local++) {
+        shared_kv[m_local*E + e] += keys[n][h][l][e] * values[n][h][l][m_start + m_local];
+        shared_kv_bw[m_local*E + e] += queries[n][h][l_b][e] * grad_out[n][h][l_b][m_start + m_local];
+        res += grad_out[n][h][l][m_start + m_local] * shared_kv[m_local*E + e];
+        res_bw += values[n][h][l_b][m_start + m_local] * shared_kv_bw[m_local*E + e];
+      }
+      atomicAdd(
+        &grad_queries[n][h][l][e],
+        res
+      );
+      atomicAdd(
+        &grad_keys[n][h][l_b][e],
+        res_bw
+      );
+    }
+}
+
+
+__global__ void causal_dot_backward_value_kernel(
+    const float_accessor queries,
+    const float_accessor keys,
+    const float_accessor values,
+    const float_accessor grad_out,
+    float_accessor grad_keys,
+    float_accessor grad_values,
+    int N,
+    int H,
+    int L,
+    int E,
+    int M
+) {
+    int n = blockIdx.y;
+    int h = blockIdx.z;
+
+    int e_start = blockIdx.x * E_BLOCK_SIZE;
+    int m = threadIdx.x % M;
+
+    extern __shared__ float shared_mem[];
+    float* shared_kv = shared_mem;
+    for (int e_local = 0; e_local < E_BLOCK_SIZE && e_local + e_start < E; e_local++) {
+      shared_kv[m + e_local * M] = 0;
+    }
+
+    for (int l = 0; l < L; l++) {
+        int l_b = L - l -1;
+        float res = 0;
+        for (int e_local = 0; e_local < E_BLOCK_SIZE && e_local + e_start < E; e_local++) {
+          shared_kv[e_local*M + m] += queries[n][h][l_b][e_start + e_local] * grad_out[n][h][l_b][m];
+          res += keys[n][h][l_b][e_start + e_local] * shared_kv[e_local*M + m];
+        }
+        atomicAdd(
+            &grad_values[n][h][l_b][m],
+            res
+        );
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void causal_dot_backward_(const torch::Tensor queries,
+                          const torch::Tensor keys,
+                          const torch::Tensor values,
+                          const torch::Tensor grad_out,
+                          torch::Tensor grad_queries,
+                          torch::Tensor grad_keys,
+                          torch::Tensor grad_values) {
+
+    // Make sure that we are using the correct GPU device
+    torch::DeviceGuard _guard(queries.device());
+
+    int N = queries.size(0);
+    int H = queries.size(1);
+    int L = queries.size(2);
+    int E = queries.size(3);
+    int M = values.size(3);
+
+    const int blocks_per_sequence = (M + M_BLOCK_SIZE - 1) / M_BLOCK_SIZE;
+
+    dim3 blockDim(E, 1, 1);
+    dim3 gridDim(blocks_per_sequence, N, H);
+    const int shared_mem_qk_backward = 2 * M_BLOCK_SIZE * E * sizeof(float);
+
+    causal_dot_backward_query_key_kernel<<<gridDim, blockDim, shared_mem_qk_backward>>>(
+      queries.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      keys.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      values.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_out.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_queries.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_keys.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      N, H, L, E, M
+    );
+
+    const int blocks_per_sequence_value = (E + E_BLOCK_SIZE - 1) / E_BLOCK_SIZE;
+
+    dim3 blockDimv(M, 1, 1);
+    dim3 gridDimv(blocks_per_sequence_value, N, H);
+    const int shared_mem_v_backward = E_BLOCK_SIZE * M * sizeof(float);
+    causal_dot_backward_value_kernel<<<gridDimv, blockDimv, shared_mem_v_backward>>>(
+      queries.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      keys.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      values.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_out.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_keys.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      grad_values.packed_accessor32<float, 4, torch::RestrictPtrTraits>(),
+      N, H, L, E, M
+    );
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void causal_dot_backward(const torch::Tensor queries,
+                         const torch::Tensor keys,
+                         const torch::Tensor values,
+                         const torch::Tensor grad_out,
+                         torch::Tensor grad_queries,
+                         torch::Tensor grad_keys,
+                         torch::Tensor grad_values) {
+#ifdef ENABLE_NVIDIA_OPTIMIZATIONS
+  int fallback = nvidia::lmha_bwd(queries,
+                                  keys,
+                                  values,
+                                  grad_out,
+                                  grad_queries,
+                                  grad_keys,
+                                  grad_values);
+#else
+  int fallback = 1;
+#endif
+  if( fallback ) {
+    // Make sure that the gradient tensors are 0. This is needed because the
+    // bwd pass might have partially executed and filled in some values in
+    // grad_queries or grad_keys.
+    //
+    // This adds a small overhead every time we have to fall back to the old
+    // kernel for the backward pass.
+    grad_queries.zero_();
+    grad_keys.zero_();
+    causal_dot_backward_(queries, keys, values, grad_out, grad_queries, grad_keys, grad_values);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def(
+        "causal_dot_product",
+        &causal_dot_product,
+        "Compute the weighted sum of values but attending only to previous "
+        "values."
+    );
+    m.def(
+        "causal_dot_backward",
+        &causal_dot_backward,
+        "Compute the gradients for the causal dot product."
+    );
+}

csrc/setup.py

@@ -0,0 +1,53 @@
+#
+# Copyright (c) 2020 Idiap Research Institute, http://www.idiap.ch/
+# Written by Angelos Katharopoulos <angelos.katharopoulos@idiap.ch>,
+# Apoorv Vyas <avyas@idiap.ch>
+#
+
+import torch
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
+import subprocess
+
+def get_last_arch_torch():
+    arch = torch.cuda.get_arch_list()[-1]
+    print(f"Found arch: {arch} from existing torch installation")
+    return arch
+
+def get_cuda_bare_metal_version(cuda_dir):
+    raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
+    output = raw_output.split()
+    release_idx = output.index("release") + 1
+    release = output[release_idx].split(".")
+    bare_metal_major = release[0]
+    bare_metal_minor = release[1][0]
+    return raw_output, bare_metal_major, bare_metal_minor
+
+def append_nvcc_threads(nvcc_extra_args):
+    _, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
+    if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:
+        return nvcc_extra_args + ["--threads", "4"]
+    return nvcc_extra_args
+
+arch = get_last_arch_torch()
+sm_num = arch[-2:]
+cc_flag = ['--generate-code=arch=compute_90,code=compute_90']    # for H100
+# cc_flag = ['--generate-code=arch=compute_80,code=compute_80']  # for A100
+# cc_flag = ['--generate-code=arch=compute_89,code=compute_89']  # for RTX 6000, 4090
+# cc_flag = ['--generate-code=arch=compute_86,code=compute_86']  # for A6000, 3090
+# cc_flag = ['--generate-code=arch=compute_75,code=compute_75']
+
+setup(
+    name='causal_attention_cuda_cpp',
+    ext_modules=[
+        CUDAExtension('causal_attention_cuda', [
+            # 'causal_attention.cpp',
+            'causal_attention_cuda.cu',
+        ],
+        extra_compile_args={'cxx': ['-O3'],
+                             'nvcc': append_nvcc_threads(['-O3', '-lineinfo', '--use_fast_math', '-std=c++17'] + cc_flag)
+                            })
+    ],
+    cmdclass={
+        'build_ext': BuildExtension
+    })

src/dataloaders/__init__.py

@@ -0,0 +1,22 @@
+"""
+Load dataloaders
+"""
+import importlib
+
+
+def load_data(dataset_config: dict, dataloader_config: dict):
+    """Return dataloaders from dataset_config"""
+    try:
+        dataset_module = importlib.import_module(f'dataloaders.{dataset_config["name"]}')
+    except Exception:
+        try:
+            dataset_module = importlib.import_module(f'src.dataloaders.{dataset_config["name"]}')
+        except Exception as e2:
+            print(e2)
+            try:  # e.g., tasks like GLUE where name is benchmark and path specifies the dataset / task
+                dataset_module = importlib.import_module(f'dataloaders.{dataset_config["path"]}')
+            except Exception as e3:
+                print(f'Error from {dataset_config}')
+                raise e3
+    _load_data = getattr(dataset_module, 'load_data')
+    return _load_data(**dataset_config, **dataloader_config)

src/dataloaders/alpaca_clean.py

@@ -0,0 +1,149 @@
+"""
+Alpaca training dataloaders
+
+We adopt the original prompt template; goes something like:
+```
+Below is an instruction that describes a task. 
+Write a response that appropriately completes the request.
+### Instruction:
+{instruction}
+ 
+### Response:
+{response}
+```
+See `PROMPT_DICT` for more. 
+"""
+from functools import partial
+from os.path import join
+
+from datasets import load_metric, load_dataset
+
+from .utils import (
+    get_lm_loader, get_seq2seq_loader,
+    convert_to_hf_dataset, 
+    get_tokenizer_from_config,
+    download_scrolls_metric as download_metric
+)
+from .utils.packing import ConcatDataset
+
+
+PROMPT_DICT = {
+    "prompt_input": (
+        "Below is an instruction that describes a task, paired with an input that provides further context. "
+        "Write a response that appropriately completes the request.\n\n"
+        "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
+    ),
+    "prompt_no_input": (
+        "Below is an instruction that describes a task. "
+        "Write a response that appropriately completes the request.\n\n"
+        "### Instruction:\n{instruction}\n\n### Response:\n"
+    ),
+}
+
+
+def load_data(name: str, dataset_config: dict, pretrained_model_config: dict,
+              preprocess_config: dict, **loader_kwargs: any):
+    """
+    Shared function to load dataset from experiment config
+    -> e.g., see configs/experiments/distill_alpaca_clean_lr1e-2.yaml
+    """
+    # Misc. setup
+    cache_dir = dataset_config['cache_dir']
+    input_len = dataset_config['chunk_size']
+    concat_data = dataset_config['concat_data']
+
+    tokenizer_name = pretrained_model_config['pretrained_model_name_or_path']
+    tokenizer_name = tokenizer_name.split('/')[-1]
+    # save_path = join(cache_dir, f'{name}_{tokenizer_name}')
+    
+    # Setup tokenizer
+    tokenizer = get_tokenizer_from_config(pretrained_model_config)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        print(f'Setting tokenizer.pad_token to {tokenizer.pad_token}')
+
+    tokenizer.padding_side = 'left'  # for decoder-only generation
+    # Get initial data
+    ignore_kwargs = ['concat_data', 'chunk_size', 'pose_kwargs']
+    dataset = load_dataset(
+        **{k: v for k, v in dataset_config.items() if k not in ignore_kwargs}
+    )
+    if dataset_config['name'] == 'samsum':  # hack
+        dataset = dataset.rename_column('dialogue', 'input')
+        dataset = dataset.rename_column('summary', 'output')
+        _instruction = 'Summarize this dialogue.'
+        for split in dataset.keys():
+            dataset[split] = dataset[split].add_column(
+                'instruction', [_instruction] * len(dataset[split])
+            )
+        train_set, val_set, test_set = dataset['train'], dataset['validation'], dataset['test']
+        dataset = train_set  # hack to work with below code
+    else:
+        dataset = dataset['train']
+        train_set = convert_to_hf_dataset([dataset[ix] for ix in range(200, len(dataset))], cache_dir)
+        val_set   = convert_to_hf_dataset([dataset[ix] for ix in range(200)], cache_dir)
+        test_set  = convert_to_hf_dataset([dataset[ix] for ix in range(200)], cache_dir)
+
+    # Convert to dicts of {input_ids, attention_mask, labels}
+    train_set = train_set.map(
+        partial(template_and_tokenize, tokenizer=tokenizer, include_label=True), 
+        remove_columns=list(dataset.features),) #  load_from_cache_file=False)
+    val_set = val_set.map(
+        partial(template_and_tokenize, tokenizer=tokenizer, include_label=True),
+        remove_columns=list(dataset.features),) #  load_from_cache_file=False)
+    test_set  = test_set.map(
+        partial(template_and_tokenize, tokenizer=tokenizer, include_label=False),
+        remove_columns=list(dataset.features),) #  load_from_cache_file=False)
+
+    # Chunk together train and val sets
+    if concat_data:
+        train_set = ConcatDataset(train_set, chunk_size=input_len)
+        val_set = ConcatDataset(val_set, chunk_size=input_len)
+
+    # Get dataloaders
+    dataloaders = {
+        'train': get_lm_loader(train_set, tokenizer, 'train', input_len, **loader_kwargs),
+        'validation': get_lm_loader(val_set, tokenizer, 'validation', input_len, **loader_kwargs),
+        'test': get_seq2seq_loader(test_set, tokenizer, 'test', **loader_kwargs),
+    }
+    # Evaluation metric
+    try:
+        metric = load_metric(download_metric(), 'gov_report')  # hack but we want rouge
+    except Exception as e:
+        print(f'Error loading metric: {e}')
+        metric = None
+
+    # Finishing touches
+    for k, v in dataloaders.items():  # Make tokenizer accessible
+        dataloaders[k].dataset.tokenizer = tokenizer
+        dataloaders[k].dataset.metric = metric
+    return dataloaders
+
+
+def template_and_tokenize(sample, tokenizer, include_label: bool = True):
+    """
+    Format dataset context and answers into single-sequence prompts
+    """
+    if sample.get('input', '') == '':
+        prompt = PROMPT_DICT["prompt_no_input"].format_map(sample)
+    else:
+        prompt = PROMPT_DICT["prompt_input"].format_map(sample)
+
+    prompt = tokenizer.encode(prompt, add_special_tokens=True)
+    if include_label:
+        answer = tokenizer.encode(f'{sample["output"]}{tokenizer.eos_token}', 
+                                  add_special_tokens=False)
+        target = None
+    else:
+        answer = []
+        target = tokenizer.encode(f'{sample["output"]}{tokenizer.eos_token}', 
+                                  add_special_tokens=False)
+    input_ids = prompt + answer
+    attn_mask = [1] * len(input_ids)
+
+    sample =  {
+        "input_ids": input_ids,
+        "attention_mask" : attn_mask,
+        "labels": [-100] * len(prompt) + answer if include_label else target,
+    }
+    return sample

src/dataloaders/alpaca_clean_instruct.py

@@ -0,0 +1,148 @@
+"""
+Alpaca Clean dataset with Llama3-Instruct prompt formatting
+"""
+
+from functools import partial
+from os.path import join
+
+import numpy as np
+from tqdm import tqdm
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+
+from datasets import load_metric, load_dataset
+from transformers import AutoTokenizer
+from transformers import DataCollatorForSeq2Seq, DefaultDataCollator, DataCollatorWithPadding
+
+from .utils import (
+    get_lm_loader, get_seq2seq_loader,
+    convert_to_hf_dataset, 
+    get_tokenizer_from_config,
+    download_scrolls_metric as download_metric
+)
+from .utils.packing import ConcatDataset
+
+
+SYSTEM_PROMPT = "You are a helpful AI assistant who always responds to appropriately complete a user's request."
+
+
+def encode_response(response: str, tokenizer) -> list[int]:
+    tokens = tokenizer.encode(response.strip(), add_special_tokens=False)
+    # For Llama 3 Instruct: tokens.append(tokenizer.get_added_vocab()["<|eot_id|>"])
+    tokens.append(tokenizer.eos_token_id)  
+    try:  # Llama 3 Instruct
+        tokens.append(tokenizer.get_added_vocab()["<|end_of_text|>"])
+    except KeyError:
+        pass
+    return tokens
+
+
+def load_data(name: str, dataset_config: dict, pretrained_model_config: dict,
+              preprocess_config: dict, **loader_kwargs: any):
+
+    # Misc. setup
+    cache_dir = dataset_config['cache_dir']
+    input_len = dataset_config['chunk_size']
+    concat_data = dataset_config['concat_data']
+    load_from_cache_file = False  # False if want to retokenize dataset
+
+    # Hard-code system prompt handling
+    if 'istral' in pretrained_model_config['pretrained_model_name_or_path']:
+        system_prompt = ''
+    else:
+        system_prompt = SYSTEM_PROMPT
+
+    tokenizer_name = pretrained_model_config['pretrained_model_name_or_path']
+    tokenizer_name = tokenizer_name.split('/')[-1]
+    save_path = join(cache_dir, f'{name}_{tokenizer_name}')
+    
+    # Setup tokenizer
+    tokenizer = get_tokenizer_from_config(pretrained_model_config)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        print(f'Setting tokenizer.pad_token to {tokenizer.pad_token}')
+
+    tokenizer.padding_side = 'left'  # for decoder-only generation
+
+    # Get initial data
+    ignore_kwargs = ['concat_data', 'chunk_size', 'pose_kwargs', 'system_prompt', 'name']
+    train_set = load_dataset(
+        **{k: v for k, v in dataset_config.items() if k not in ignore_kwargs},
+        split='train[100:-100]',
+    )
+    val_set = load_dataset(  # we just use this dataset as a validation set
+        **{k: v for k, v in dataset_config.items() if k not in ignore_kwargs},
+        split='train[:100]+train[-100:]',
+    )
+    test_set = load_dataset(
+        **{k: v for k, v in dataset_config.items() if k not in ignore_kwargs},
+        split='train[:100]+train[-100:]',
+    )
+
+    # Convert to dicts of {input_ids, attention_mask, labels}
+    train_set = train_set.map(partial(template_and_tokenize, tokenizer=tokenizer, 
+                                      include_label=True, system_prompt=system_prompt),
+                              remove_columns=list(train_set.features), 
+                              load_from_cache_file=load_from_cache_file)
+    val_set   = val_set.map(partial(template_and_tokenize, tokenizer=tokenizer, 
+                                    include_label=True, system_prompt=system_prompt),
+                            remove_columns=list(val_set.features),
+                            load_from_cache_file=load_from_cache_file)
+    test_set  = test_set.map(partial(template_and_tokenize, tokenizer=tokenizer, 
+                                     include_label=False, system_prompt=system_prompt),
+                             remove_columns=list(test_set.features),
+                             load_from_cache_file=load_from_cache_file)
+
+    # Chunk together train and val sets
+    if concat_data:
+        train_set = ConcatDataset(train_set, chunk_size=input_len)
+        val_set = ConcatDataset(val_set, chunk_size=input_len)
+    
+    # Get dataloaders
+    dataloaders = {
+        'train': get_lm_loader(train_set, tokenizer, 'train', input_len, **loader_kwargs),
+        'validation': get_lm_loader(val_set, tokenizer, 'validation', input_len, **loader_kwargs),
+        'test': get_seq2seq_loader(test_set, tokenizer, 'test', **loader_kwargs),
+    }
+    # Evaluation metric
+    metric = load_metric(download_metric(), 'gov_report')  # hack but we want rouge
+    
+    # Finishing touches
+    for k, v in dataloaders.items():  # Make tokenizer accessible
+        dataloaders[k].dataset.tokenizer = tokenizer
+        dataloaders[k].dataset.metric = metric
+    return dataloaders
+
+
+def template_and_tokenize(sample, tokenizer, include_label: bool = True, 
+                          system_prompt: str = None):
+    if system_prompt is None:
+        system_prompt = SYSTEM_PROMPT
+
+    prompt = sample['instruction']
+    if sample['input'] != '':
+        prompt += f"\n\n{sample['input']}"
+    
+    messages = [
+        {"role": "system", "content": system_prompt},
+    ] if system_prompt != '' else []
+    messages.append({"role": "user", "content": prompt})
+    prompt_ids = tokenizer.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True,
+    )
+    if include_label:
+        answer = encode_response(sample['output'], tokenizer)
+    else:
+        answer = []
+        target = encode_response(sample['output'], tokenizer)
+        
+    input_ids = prompt_ids + answer
+    attn_mask = [1] * len(input_ids)
+    sample =  {
+        "input_ids": input_ids,
+        "attention_mask" : attn_mask,
+        "labels": [-100] * len(prompt_ids) + answer if include_label else target,
+    }
+    return sample
+

src/dataloaders/utils/__init__.py

@@ -0,0 +1,4 @@
+"""
+Helper functions dataset setup and loading
+"""
+from .setup import *

src/dataloaders/utils/llama3.py

@@ -0,0 +1,62 @@
+"""
+Data utils for Llama3
+"""
+
+def encode_header(message: str, tokenizer) -> list[int]:
+    tokens = []
+    tokens.append(tokenizer.get_added_vocab()["<|start_header_id|>"])
+    tokens.extend(tokenizer.encode(message["role"], add_special_tokens=False))
+    tokens.append(tokenizer.get_added_vocab()["<|end_header_id|>"])
+    tokens.extend(tokenizer.encode("\n\n", add_special_tokens=False))
+    return tokens
+
+
+def encode_message(message: str, tokenizer, include_header: bool = True) -> list[int]:
+    tokens = encode_header(message, tokenizer) if include_header else []
+    tokens.extend(
+        tokenizer.encode(message["content"].strip(), add_special_tokens=False)
+    )
+    tokens.append(tokenizer.get_added_vocab()["<|eot_id|>"])
+    return tokens
+
+
+def template_and_tokenize(sample, tokenizer, include_label: bool = True, 
+                          system_prompt: str = None):
+    if system_prompt is not None:
+        dialog = [{'role': 'system', 'content': system_prompt}]
+    else:
+        dialog = []
+
+    chat = []
+    instruction = sample['instruction']
+    if sample['input'] != '':
+        instruction += f"\n\n{sample['input']}"
+    dialog.extend([
+        {'role': 'user', 'content': instruction},
+        {'role': 'assistant', 'content': sample['output']},
+    ])
+        
+    prompt = []
+    prompt.append(tokenizer.get_added_vocab()["<|begin_of_text|>"])
+    for message in dialog[:-1]:
+        prompt.extend(encode_message(message, tokenizer))
+
+    if include_label:
+        answer = encode_message(dialog[-1], tokenizer)
+        answer.append(tokenizer.get_added_vocab()["<|end_of_text|>"])
+    else:
+        answer = []
+        target = encode_message(dialog[-1], tokenizer, include_header=False)
+        target.append(tokenizer.get_added_vocab()["<|end_of_text|>"])
+        # Add the start of an assistant message for the model to complete.
+        prompt.extend(encode_header({"role": "assistant", "content": ""}, tokenizer))
+
+    input_ids = prompt + answer
+    attn_mask = [1] * len(input_ids)
+
+    sample =  {
+        "input_ids": input_ids,
+        "attention_mask" : attn_mask,
+        "labels": [-100] * len(prompt) + answer if include_label else target,
+    }
+    return sample

src/dataloaders/utils/packing.py

@@ -0,0 +1,80 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.
+"""
+Copied from https://github.com/meta-llama/llama-recipes/blob/9b3dabcaac78980eae40005bbc8b1a8276c82af3/src/llama_recipes/data/concatenator.py#L1
+"""
+import random
+from itertools import chain
+from tqdm import tqdm
+
+
+from torch.utils.data import Dataset
+
+
+class Concatenator(object):
+    def __init__(self, chunk_size=2048):
+        self.chunk_size=chunk_size
+        self.residual = {"input_ids": [], "attention_mask": []}
+        
+    def __call__(self, batch):
+        concatenated_samples = {
+            k: v + list(chain(*batch[k])) for k, v in self.residual.items()
+        }
+
+        total_length = len(concatenated_samples[list(concatenated_samples.keys())[0]])
+
+        if total_length >= self.chunk_size:
+            chunk_num = total_length // self.chunk_size
+            result = {
+                k: [
+                    v[i : i + self.chunk_size]
+                    for i in range(0, chunk_num * self.chunk_size, self.chunk_size)
+                ]
+                for k, v in concatenated_samples.items()
+            }
+            self.residual = {
+                k: v[(chunk_num * self.chunk_size) :]
+                for k, v in concatenated_samples.items()
+            }
+        else:
+            result = concatenated_samples
+            self.residual = {k: [] for k in concatenated_samples.keys()}
+
+        result["labels"] = result["input_ids"].copy()
+
+        return result
+
+class ConcatDataset(Dataset):
+    """
+    Concatenates or packs samples of a dataset into chunks of size `chunk_size`
+    """
+    def __init__(self, dataset, chunk_size: int = 1024, seed: int = 42,) -> None:
+        self.dataset = dataset
+        self.chunk_size = chunk_size
+        self.samples = []
+        buffer = {
+            "input_ids": [],
+            "attention_mask": [],
+            "labels": [],
+            }
+        random.seed(seed)
+        for sample in tqdm(self.dataset, desc="Preprocessing dataset", dynamic_ncols=True):
+            buffer = {k: v + sample[k] for k,v in buffer.items()}
+            
+            while len(next(iter(buffer.values()))) > self.chunk_size:
+                self.samples.append({k: v[:self.chunk_size] for k,v in buffer.items()})
+                buffer = {k: v[self.chunk_size:] for k,v in buffer.items()}
+        # Slow hack, but filter out any samples without valid labels (all -100)
+        self.filtered_samples = []
+        for s in self.samples:
+            if sum(s['labels']) != chunk_size * -100:
+                self.filtered_samples.append(s)
+        if len(self.filtered_samples) < len(self.samples):
+            print(f'OG dataset: {len(self.samples)} samples -> Filtered dataset: {len(self.filtered_samples)}')
+            print(f'-> Filtered out {len(self.samples) - len(self.filtered_samples)} samples')
+                
+    def __getitem__(self, idx):
+        return self.filtered_samples[idx]
+    
+    def __len__(self):
+        return len(self.filtered_samples)

src/dataloaders/utils/setup.py

@@ -0,0 +1,123 @@
+"""
+Helper functions dataset setup and loading
+"""
+import os
+from os.path import join
+import shutil
+import numpy as np
+
+from torch.utils.data import Dataset, DataLoader
+
+from datasets import Dataset as HFDataset
+from huggingface_hub import hf_hub_download
+from transformers import AutoTokenizer, LlamaTokenizer
+from transformers import DataCollatorForSeq2Seq  
+# from transformers import DefaultDataCollator, DataCollatorWithPadding
+
+
+def get_seq2seq_loader(dataset: Dataset, tokenizer: AutoTokenizer,
+                       split: str, **loader_kwargs: any):
+    """
+    Get dataloader for seq2seq tasks (evaluation)
+    """
+    tokenizer.padding_side = 'right'
+    collate_fn = DataCollatorForSeq2Seq(
+        tokenizer, label_pad_token_id=-100, return_tensors='pt')
+    return DataLoader(
+        dataset, shuffle='train' in split, collate_fn=collate_fn, **loader_kwargs)
+
+
+def get_lm_loader(dataset: Dataset, tokenizer: AutoTokenizer,
+                  split: str, max_length: int = None, **loader_kwargs: any):
+    """
+    Get dataloader for language modeling (training)
+    -> Currently this ends up being the same as get_seq2seq_loader
+    """
+    # collate_fn = DefaultDataCollator(return_tensors='pt')
+    # collate_fn = DataCollatorWithPadding(tokenizer=tokenizer, padding=True,
+    #                                      max_length=max_length, return_tensors='pt')
+    collate_fn = DataCollatorForSeq2Seq(
+        tokenizer, label_pad_token_id=-100, return_tensors='pt')
+    return DataLoader(
+        dataset, shuffle='train' in split, collate_fn=collate_fn, **loader_kwargs)
+
+
+def convert_to_hf_dataset(dataset, cache_dir: str):
+    """
+    Convert iterable dataset to HuggingFace HFDataset object
+    """
+    def gen():
+        for _, sample in enumerate(dataset):
+            yield sample  # dataset[idx]
+    return HFDataset.from_generator(gen, cache_dir=cache_dir)
+
+
+def get_tokenizer_from_config(model_config):
+    """
+    Get pretrained tokenizer based on (pretrained) model config
+    """
+    # Get tokenizer
+    if 'llama' in model_config['pretrained_model_name_or_path']:
+        try:  # if we store locally
+            model_path = join(model_config['cache_dir'],
+                              model_config['pretrained_model_name_or_path'])
+            tokenizer = LlamaTokenizer.from_pretrained(model_path)
+        except Exception as e:
+            try:
+                tokenizer = AutoTokenizer.from_pretrained(**model_config)
+                print("-> Bad LlamaTokenizer.from_pretrained(model_path)", e)
+                print("-> But resolved with: AutoTokenizer.from_pretrained(**model_config)")
+            except Exception as e2:
+                print("-> Error with AutoTokenizer.from_pretrained(**model_config)", e2)
+            # tokenizer = LlamaTokenizer.from_pretrained(**model_config)  # v4.43 errors with `*** TypeError: not a string`
+    elif 'Mistral-7B-Instruct-v0.3' in model_config['pretrained_model_name_or_path']:
+        tokenizer = LlamaTokenizer.from_pretrained(**model_config)  # hack where AutoTokenizer doesn't recognize
+    elif 'Mistral-7B' in model_config['pretrained_model_name_or_path']:
+        tokenizer = AutoTokenizer.from_pretrained(**model_config)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(**model_config)
+    return tokenizer
+
+
+def add_special_tokens_to_dataset(dataset, tokenizer):
+    """
+    Add special tokens as attributes to a dataset object
+    """
+    token_map = {k: v for k, v in tokenizer.special_tokens_map.items()}
+    special_ids = tokenizer.all_special_ids
+    for idx, k in enumerate(tokenizer.special_tokens_map.keys()):
+        token_map[f'{k}_id'] = special_ids[idx]
+    for k, v in token_map.items():
+        setattr(dataset, k, v)
+    return dataset
+
+
+def train_test_split(samples: any, train_size: int, test_size: int, seed: int):
+    """
+    Split samples into train and test sets
+    """
+    try:
+        assert len(samples) == train_size + test_size
+    except Exception as e:
+        print(len(samples), train_size + test_size)
+        raise e
+    arange = np.arange(len(samples))
+    np.random.seed(seed)
+    test_idx = np.random.choice(arange, size=test_size, replace=False)
+    train_idx = np.setdiff1d(arange, test_idx)
+    return samples[train_idx], samples[test_idx]
+
+
+def download_scrolls_metric():
+    """
+    Download ROUGE, F1, and other accuracy metrics included in the SCROLLS dataset
+    """
+    scrolls_metric_path = hf_hub_download(
+        repo_id="tau/scrolls", filename="metrics/scrolls.py", repo_type="dataset"
+    )
+    updated_scrolls_metric_path = (
+        os.path.dirname(scrolls_metric_path) + 
+        os.path.basename(scrolls_metric_path).replace(".", "_") + ".py"
+    )
+    shutil.copy(scrolls_metric_path, updated_scrolls_metric_path)
+    return updated_scrolls_metric_path

src/finetune.py

@@ -0,0 +1,68 @@
+"""
+Finetuning functions to do post-distillation
+"""
+from os.path import join
+from omegaconf import OmegaConf
+
+import torch
+from torch.nn import Module
+
+from src.utils.setup import update_config_from_args
+from src.dataloaders import load_data
+from src.trainer import get_trainer, get_optimizer, get_scheduler
+
+
+def prepare_finetune_configs(args, model_config: dict,
+                             finetune_config_name: str = None,
+                             finetune_checkpoint_name: str = None,
+                             config_dir='./configs/experiment'):
+    """
+    Prepare finetuning configs
+    """
+    # Load finetuning config
+    finetune_config = (finetune_config_name if finetune_config_name is not None else 
+                       finetune_checkpoint_name.split('-f=')[-1].split('-')[0])
+    finetune_config_path = join(config_dir, f'{finetune_config}.yaml')
+    finetune_config = OmegaConf.load(finetune_config_path)
+    finetune_config = update_config_from_args(finetune_config, args,
+                                              ignore_args=['lr', 'weight_decay'])
+    # Update data tokenizer to match model
+    if getattr(finetune_config.dataset, 'pretrained_model_config', None) is not None:
+        for k in ['pretrained_model_name_or_path', 'cache_dir']:
+            finetune_config.dataset.pretrained_model_config[k] = model_config['model'][k]
+    # Set finetuning args
+    for arg, argv in finetune_config.trainer.items():
+        if arg != 'name':
+            setattr(args, arg, argv)
+    for _config in ['dataloader', 'optimizer', 'lr_scheduler']:
+        setattr(args, _config, OmegaConf.to_container(getattr(finetune_config, _config)))
+    return finetune_config, args
+
+
+def get_finetuner(model: Module, finetune_config: dict, device: torch.device, 
+                  args: any, wandb: any, initial_eval: bool = False):
+    """
+    Initialize finetuning trainer
+    """
+    model.to(device)  # if using a fused optimizer
+    model.train()
+
+    # Initialize optimizer and scheduler
+    optimizer = get_optimizer(model=model, **finetune_config.optimizer)
+    scheduler = get_scheduler(optimizer=optimizer, **finetune_config.lr_scheduler)
+
+    dataloaders  = load_data(finetune_config.dataset, finetune_config.dataloader)
+    train_loader = dataloaders[finetune_config.trainer.train_split]
+    eval_loader  = dataloaders[finetune_config.trainer.val_split]
+
+    OurTrainer = get_trainer(finetune_config.trainer.name)
+    trainer = OurTrainer(model=model,
+                         args=args,
+                         train_loader=train_loader,
+                         eval_loader=eval_loader,
+                         optimizer_and_scheduler=(optimizer, scheduler),
+                         device=device,
+                         wandb=wandb,
+                         checkpoint_suffix='_ft',
+                         **finetune_config.trainer)
+    return trainer

src/model/convert_model.py

@@ -0,0 +1,173 @@
+"""
+Attention conversion helpers
+"""
+from functools import partial
+from tqdm import tqdm
+import torch.nn as nn
+
+
+def convert_attention(model: nn.Module, 
+                      attention_config: dict, 
+                      train_attention: bool = False,
+                      remove_base_attn: bool = True,):
+    """
+    Call to convert all attention layers
+    """
+    softmax_attns = []
+    if 'softmax_attentions' in attention_config:
+        softmax_attns = attention_config['softmax_attentions']
+    if attention_config.attention_type != 'softmax':
+        layers = traverse_layers(model)
+        for layer_idx, layer in enumerate(tqdm(layers, desc='Converting attentions...')):
+            if layer_idx not in softmax_attns:
+                layer.self_attn = convert_llama_attention(
+                    layer, attention_config, layers, train_attention, remove_base_attn,
+                )
+                layer.self_attn.converted = True
+            else:  # Freeze any preserved softmax attention layers
+                for p in layer.parameters():
+                    p.requires_grad = False
+    else:
+        print(f'-> attention_config.attention_type is {attention_config.attention_type}; not converting attentions')
+    return model
+
+
+def toggle_attention(llama_model: nn.Module, train: bool = False):
+    """
+    Make attentions trainable if train is True
+    -> Set train_attention = False when finetuning
+    """
+    for layer in traverse_layers(llama_model):
+        layer.self_attn.train_attention = train
+    return llama_model
+
+
+def remove_base_attention(llama_model: nn.Module):
+    """
+    Remove teacher attention after distillation (if we keep it)
+    """
+    for layer in traverse_layers(llama_model):
+        if getattr(layer.self_attn, 'base_attn', False):
+            del layer.self_attn.base_attn
+    return llama_model
+        
+
+def traverse_layers(model: nn.Module, verbose: bool = False):
+    """
+    Return list of model layers
+    """
+    try:
+        layers = model.model.layers
+        if verbose:
+            print('-> Loading from model.model.layers')
+    except AttributeError as e: # if base model
+        if verbose:
+            print(e)
+        try:
+            layers = model.layers
+            if verbose:
+                print('-> Loading from model.layers')
+        except AttributeError as e1:  # If we make a PEFT model
+            if verbose:
+                print(e1)
+            layers = model.base_model.model.model.layers
+            if verbose:
+                print('-> Loading from model.base_model.model.model.layers')
+    return layers
+
+
+def convert_llama_attention(layer: nn.Module,
+                            attention_config: dict,
+                            layers: list[nn.Module],  # list of layers
+                            train_attention: bool = False,
+                            remove_base_attn: bool = True):
+    """
+    Converts a single layer's attention layer as specified by attention_config
+    """
+    return get_attention(**attention_config)(
+        base_attn=layer.self_attn,
+        layer_idx=layer.self_attn.layer_idx,  # Transformers v4.36
+        max_layer_idx=len(layers) - 1,
+        train_attention=train_attention,
+        remove_base_attn=remove_base_attn,
+    )
+
+
+def get_attention(attention_type: str, **kwargs: any):
+    """
+    Get the linear attention class; either purely linear or linear with sliding window
+    -> 'linear' == 'lolcats_llama'
+    -> 'linear and sliding_window' == 'lolcats_llama_window_*'
+    """
+    kwargs['attention_type'] = attention_type
+
+    if attention_type == 'lolcats_llama':
+        from .linear_attention import LolcatsLinearAttention
+        return partial(LolcatsLinearAttention, **kwargs)
+
+    elif attention_type == 'lolcats_llama_window_tk':
+        from .linear_attention import LolcatsTKWindowAttention
+        return partial(LolcatsTKWindowAttention, **kwargs)
+
+    elif attention_type == 'lolcats_llama_window_sw':
+        from .linear_attention import LolcatsSlidingWindowAttention
+        return partial(LolcatsSlidingWindowAttention, **kwargs)
+
+    elif attention_type == 'lolcats_llama_window_sw_linear':
+        from .linear_attention.linear_window_attention_sw_linear import LolcatsLinearSlidingWindowAttention
+        return partial(LolcatsLinearSlidingWindowAttention, **kwargs)
+
+    ## Experimental chunked linear attentions below
+    elif attention_type == 'lolcats_long_llama_window_tk':
+        from .linear_attention import LolcatsTKWindowLongAttention
+        return partial(LolcatsTKWindowLongAttention, **kwargs)
+
+    elif attention_type == 'lolcats_long_llama_window_sw':
+        from .linear_attention import LolcatsSlidingWindowLongAttention
+        return partial(LolcatsSlidingWindowLongAttention, **kwargs)
+
+    ## TK generation build (requires Thunderkittens)
+    elif attention_type == 'lolcats_llama_window_tk_gen':
+        from .linear_attention import LolcatsWindowAttentionTKGen
+        return partial(LolcatsWindowAttentionTKGen, **kwargs)
+
+    else:
+        print(f'-> attention_type {attention_type} not handled... returning None')
+        return None
+
+
+def get_attention_cache(attention_type: str, past_key_values: any = None):
+    """
+    Determine how we store past keys and values when generating
+    """
+    if attention_type is None:
+        return past_key_values
+
+    # print(f'Returning attention cache based on attention_type == {attention_type}')
+    elif 'lolcats_llama_window_tk_gen' in attention_type:
+        from .linear_attention import LinearAttentionTKWindowGenerationCache
+        return LinearAttentionTKWindowGenerationCache()
+
+    elif 'llama_window_tk' in attention_type:
+        from .linear_attention import LinearAttentionTKWindowCache
+        return LinearAttentionTKWindowCache()
+
+    elif 'llama_window_sw' in attention_type:
+        from .linear_attention import LinearAttentionSlidingWindowCache
+        return LinearAttentionSlidingWindowCache()
+
+    elif 'llama_window_sw_linear' in attention_type:
+        from .linear_attention import LinearAttentionSlidingWindowCache
+        return LinearAttentionSlidingWindowCache()
+
+    ## TK generation build (requires Thunderkittens)
+    elif attention_type == 'lolcats_llama_window_tk_gen':
+        from .linear_attention.linear_window_attention_tk_gen import LinearAttentionTKWindowGenerationCache
+        return LinearAttentionTKWindowGenerationCache()
+
+    elif 'softmax' in attention_type:
+        return past_key_values
+
+    else:
+        from .linear_attention import LinearAttentionState
+        return LinearAttentionState()

src/model/feature_map.py

@@ -0,0 +1,306 @@
+"""
+Learnable linear attention feature map classes and functions
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def init_feature_map(name: str, mlp: nn.Module, **kwargs: dict):
+    """
+    Initialize feature map final activation for linear attention
+    """
+    return FeatureMap(activation_name=name, mlp=mlp, **kwargs)
+
+
+def init_feature_map_act(name: str, fullspace: bool = True, **kwargs):
+    """
+    Initialize feature map final activation for linear attention
+    """
+    if name == 'softmax_dim' and fullspace:
+        return SoftmaxDim(**kwargs)
+    elif name == 'softmax_dim' and not fullspace:
+        return SoftmaxDimHalfspace(**kwargs)
+    elif name == 'exp_dim' and fullspace:
+        return Exp(**kwargs)
+    elif name == 'exp_dim' and not fullspace:
+        return ExpHalfspace(**kwargs)
+    elif name == 'pos_elu':
+        return PosELU(**kwargs)
+    elif name == 'relu':
+        return ReLU(**kwargs)
+    
+    else:
+        raise NotImplementedError
+
+
+def init_learned_kernel(name: str, **kwargs: any):
+    """
+    Initialize feature map MLP for linear attention
+    """
+    if name == 'untied_head_einsum':
+        return FeatureMapMLP(**kwargs)
+    elif name == 'untied_head_adapter':
+        return FeatureMapAdapter(**kwargs)
+    else:
+        raise NotImplementedError
+
+
+class FeatureMap(nn.Module):
+    """
+    Final 'activation' of feature map. Can probably be combined with
+    `FeatureMapMLP` below
+
+    Full feature map is like f(xW + b)
+    -> This is the `f` part
+    """
+    def __init__(self,
+                 activation_name: str,
+                 head_dim_idx: int = -1, 
+                 eps: float = 1e-12, 
+                 mlp: nn.Module = None,
+                 fullspace: bool = True,):
+        super().__init__()
+        self.head_dim_idx = head_dim_idx     
+        self.eps = eps
+        self.mlp = mlp if mlp is not None else nn.Identity()
+        self.activation = init_feature_map_act(activation_name, fullspace, eps=eps)
+        
+    def forward(self, x: torch.Tensor, *mlp_args: any, **mlp_kwargs: any):
+        """
+        Assume x.shape is (batch_size, n_heads, seq_len, head_dim)
+        """
+        return self.activation(self.mlp(x, *mlp_args, **mlp_kwargs), x)
+
+    def q_map(self, *args: any, **kwargs: any):
+        """
+        Use for inference in case q and k feature maps differ
+        """
+        return self.forward(*args, **kwargs)
+
+    def k_map(self, *args: any, **kwargs: any):
+        """
+        Use for inference in case q and k feature maps differ
+        """
+        return self.forward(*args, **kwargs)
+
+
+# -----------------------
+# Feature map activations
+# -----------------------
+class FeatureMapAct(nn.Module):
+    """
+    Base class for feature map activations
+    """
+    def __init__(self, eps: float = 1e-12):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor, *args: any, **kwargs: any):
+        """
+        x.shape is (batch_size, n_heads, seq_len, head_dim)
+        """
+        return x
+
+
+class PosELU(FeatureMapAct):
+    """
+    1 + ELU activation as in https://arxiv.org/abs/2006.16236
+    """
+    def forward(self, x: torch.Tensor, *args: any, **kwargs: any):
+        return (1 + F.elu(x)).clamp(min=self.eps)
+
+
+class ReLU(FeatureMapAct):
+    """
+    ReLU activation as in https://arxiv.org/abs/2103.13076
+    """
+    def forward(self, x: torch.Tensor, *args: any, **kwargs: any):
+        return F.relu(x).clamp(min=self.eps)
+
+
+class SoftmaxDim(FeatureMapAct):
+    """
+    Softmax activation as in https://arxiv.org/abs/2402.04347
+    """
+    def forward(self, x: torch.Tensor, *args: any, **kwargs: any):
+        return torch.cat([
+            torch.softmax(x, dim=-1), torch.softmax(-x, dim=-1)
+        ], dim=-1).clamp(min=self.eps)
+
+                         
+class SoftmaxDimHalfspace(FeatureMapAct):
+    """
+    Softmax activation as in https://arxiv.org/abs/2402.04347
+    """
+    def forward(self, x: torch.Tensor, *args: any, **kwargs: any):
+        return torch.softmax(x, dim=-1).clamp(min=self.eps)
+
+
+class Exp(FeatureMapAct):
+    """
+    Exp activation as in https://arxiv.org/abs/2402.04347
+    """
+    def forward(self, x: torch.Tensor, *args: any, **kwargs: any):
+        x_max = torch.amax(x, dim=-1, keepdim=True)
+        x_min = torch.amin(x, dim=-1, keepdim=True)
+        return torch.cat([
+            torch.exp(x - x_max), torch.exp(-x + x_min)
+        ], dim=-1).clamp(min=self.eps)
+
+
+class ExpHalfspace(FeatureMapAct):
+    """
+    Exp activation as in https://arxiv.org/abs/2402.04347
+    """
+    def forward(self, x: torch.Tensor, *args: any, **kwargs: any):
+        x_max = torch.amax(x, dim=-1, keepdim=True)
+        return torch.exp(x - x_max).clamp(min=self.eps)
+
+
+# ----------------
+# Feature map MLPs
+# ----------------
+
+class FeatureMapMLP(nn.Module):
+    """
+    Learnable MLP in feature map.
+    
+    Full feature map is like f(xW + b)
+    -> This is the `W` and (optional) `b` part
+    """
+    def __init__(self, 
+                 num_heads: int,
+                 head_dim: int,     # input dim
+                 feature_dim: int,  # output dim
+                 dtype: torch.dtype,
+                 device: torch.device,
+                 skip_connection: bool = False,
+                 bias: bool = False,
+                 zero_init: bool = False,
+                 normal_init: bool = False,):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.feature_dim = feature_dim
+        self.dtype = dtype
+        self.device = device
+        self.skip_connection = skip_connection
+        self.bias = bias
+        self.zero_init = zero_init
+        self.normal_init = normal_init
+        self.init_weights_()
+
+        if self.zero_init:  # Zero-out weights or set as identity post-initialization
+            self.zero_init_with_skip_() if self.skip_connection else self.zero_init_()
+
+        if self.normal_init:
+            with torch.no_grad():
+                nn.init.normal_(self.layer)
+        
+        if self.skip_connection:
+            assertion_fail = f'If self.skip_connection we need self.head_dim == self.feature_dim but self.head_dim is {self.head_dim} != self.feature_dim is {self.feature_dim}'
+            assert self.head_dim == self.feature_dim, assertion_fail
+
+    def init_weights_(self):
+        """
+        Initialize (W)eights and (b)iases
+        """
+        self.layer = nn.Parameter(torch.zeros(
+            (self.num_heads, self.head_dim, self.feature_dim),
+            dtype=self.dtype, device=self.device,
+        ))
+        nn.init.kaiming_uniform_(self.layer)
+
+        if self.bias:
+            self.bias = nn.Parameter(torch.zeros(
+                (1, self.num_heads, 1, 1),  # self.feature_dim),
+                dtype=self.dtype, device=self.device,
+            ))
+            nn.init.kaiming_uniform_(self.bias)
+        else:
+            self.bias = 0.  # hack
+
+    def zero_init_with_skip_(self):
+        """
+        Initialize weights to zero matrix if skip connection
+        """
+        with torch.no_grad():
+            nn.init.zeros_(self.layer)
+
+    def zero_init_(self):
+        """
+        Initialize weights to identity matrix if no skip connection
+        """
+        with torch.no_grad():
+            for i in range(self.layer.shape[0]):
+                try:
+                    nn.init.eye_(self.layer[i])
+                except RuntimeError:
+                    with torch.no_grad():
+                        dtype = self.layer[i].dtype
+                        weight = torch.eye(*self.layer[i].shape,
+                                           requires_grad=self.layer[i].requires_grad,
+                                           device=self.layer[i].device)
+                        self.layer[i] = weight.to(dtype=dtype)
+
+    def forward(self, x: torch.Tensor):
+        """
+        Assume x.shape is (batch_size, num_heads, seq_len, head_dim)
+        """
+        _x = torch.einsum('hdf,bhld->bhlf', self.layer, x) + self.bias
+        return x + _x if self.skip_connection else _x
+
+
+class FeatureMapAdapter(FeatureMapMLP):
+    """
+    Learnable Feature map with bottleneck adapter
+    as in https://arxiv.org/abs/1902.00751
+
+    We don't use but could be fun to try
+    """
+    def __init__(self, hidden_dim: int, *args, **kwargs):
+        kwargs['skip_connection'] = True
+        kwargs['bias'] = True
+        kwargs['zero_init'] = True
+        self.hidden_dim = hidden_dim
+        super().__init__(*args, **kwargs)
+    
+    def init_weights_(self):
+        """
+        Initialize (W)eights and (b)iases
+        """
+        kwargs = {'dtype': self.dtype, 'device': self.device}
+        self.layer0 = nn.Parameter(
+            torch.zeros((self.num_heads, self.head_dim, self.hidden_dim), **kwargs)
+        )
+        self.layer1 = nn.Parameter(
+            torch.zeros((self.num_heads, self.hidden_dim, self.feature_dim), **kwargs)
+        )
+        nn.init.kaiming_uniform_(self.layer0)
+        nn.init.kaiming_uniform_(self.layer1)
+
+        self.bias0 = nn.Parameter(torch.zeros((1, self.num_heads, 1, self.hidden_dim), **kwargs))
+        self.bias1 = nn.Parameter(torch.zeros((1, self.num_heads, 1, self.feature_dim), **kwargs))
+        nn.init.kaiming_uniform_(self.bias0)
+        nn.init.kaiming_uniform_(self.bias1)
+
+    def zero_init_with_skip_(self):
+        with torch.no_grad():
+            nn.init.zeros_(self.layer0)
+            nn.init.zeros_(self.layer1)
+            nn.init.zeros_(self.bias0)
+            nn.init.zeros_(self.bias1)
+
+    def zero_init_(self):
+        assert NotImplementedError
+
+    def forward(self, x: torch.Tensor):
+        """
+        Assume x.shape is (batch_size, num_heads, seq_len, head_dim)
+        -> Down-project, apply nonlinearity, up-project; add skip connection
+        """
+        _x = torch.einsum('hde,bhld->bhle', self.layer0, x) + self.bias0
+        _x = F.relu(_x)
+        _x = torch.einsum('hef,bhle->bhlf', self.layer1, _x) + self.bias1
+        return x + _x if self.skip_connection else _x

src/model/linear_attention/__init__.py

@@ -0,0 +1,23 @@
+"""
+Linear and linear attention + sliding window classes
+"""
+from .linear_attention import (
+    LolcatsLinearAttention, LinearAttentionState
+)
+from .linear_window_attention_tk import (
+    LolcatsTKWindowAttention, LinearAttentionTKWindowCache
+)
+from .linear_window_attention_sw import (
+    LolcatsSlidingWindowAttention, LinearAttentionSlidingWindowCache
+)
+# Experimental chunk linear attentions
+from .linear_window_attention_tk_long import (
+    LolcatsTKWindowLongAttention,
+)
+from .linear_window_attention_sw_long import (
+    LolcatsSlidingWindowLongAttention,
+)
+from .linear_window_attention_tk_gen import (
+    LolcatsWindowAttentionTKGen,
+    LinearAttentionTKWindowGenerationCache
+)

src/model/linear_attention/linear_attention.py

@@ -0,0 +1,459 @@
+"""
+Linear attention classes
+"""
+from typing import List, Tuple, Optional
+import copy
+import torch
+import torch.nn as nn
+from omegaconf import OmegaConf, DictConfig
+
+from transformers.cache_utils import Cache  # starting at Transformers v4.36
+
+# Causal linear attention dot product CUDA kernel from fast-transformers
+try:
+    from csrc import causal_dot_product as fast_causal_dot_product
+except ImportError:
+    fast_causal_dot_product = None
+
+from src.model.feature_map import init_feature_map, init_learned_kernel
+from src.model.rotary import get_rotary_embeddings, apply_rotary_pos_emb
+from .utils import repeat_kv
+
+
+# -------------------
+# Attention functions
+# -------------------
+
+def causal_dot_product(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
+    """
+    Causal linear attention dot product
+    - If available, use CUDA kernel from fast-transformers
+    """
+    if fast_causal_dot_product is None:
+        kv = torch.einsum('bhlf,bhld->bhlfd', k, v)
+        return torch.einsum('bhlf,bhlfd->bhld', q, kv.cumsum(dim=2))
+    return fast_causal_dot_product(q, k, v)
+
+def linear_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+                     fp32_attention: bool = False, eps: float = 1e-12,
+                     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    """
+    Compute linear attention with CUDA kernel implementation from fast-transformers
+    - https://github.com/idiap/fast-transformers
+    - Assume q, k are shape (batch_size, num_heads, seq_len, feature_dim); 
+      v is shape (b, h, l, head_dim)
+    """
+    dtype = q.dtype
+    # Causal mask already applied
+    y = causal_dot_product(q.contiguous().to(dtype=torch.float32),
+                           k.contiguous().to(dtype=torch.float32),
+                           v.contiguous().to(dtype=torch.float32))
+    if fp32_attention:
+        y = (y / (torch.einsum(
+            "bhld,bhld->bhl", q.float(), k.float().cumsum(dim=2)
+        ) + eps)[..., None]).to(dtype=dtype)
+    else:
+        y = y.to(dtype=dtype)
+        k = k.float().cumsum(dim=2).to(dtype=dtype)
+        y = y / (torch.einsum("bhld,bhld->bhl", q, k) + eps)[..., None]
+    return y, None, None
+
+
+def softmax_attention(q: torch.Tensor, k: torch.Tensor, v: Optional[torch.Tensor] = None, 
+                      causal: bool = True, fp32_attention: bool = True,
+                      ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    """
+    Standard softmax attention; only compute outputs if v is not None
+    -> Assume q, k, v are shape (batch_size, num_heads, seq_len, head_dim)
+    """
+    y = None
+    a = torch.einsum('bhmd,bhnd->bhmn', q, k) * (k.shape[-1] ** -0.5)
+    if causal:  # Apply causal mask
+        m, n = a.shape[-2:]
+        causal_mask = torch.ones((m, n), device = a.device, dtype = torch.bool).triu(n - m + 1)
+        a = a.masked_fill(causal_mask, -torch.finfo(a.dtype).max)
+    if fp32_attention:
+        a = torch.softmax(a, dim=-1, dtype=torch.float32).to(q.dtype)
+    else:
+        a = torch.softmax(a, dim=-1)
+    if v is not None:
+        y = torch.einsum('bhmn,bhnd->bhmd', a, v)
+    return y, a, None
+
+
+def quadratic_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor = None,
+                        causal: bool = True, fp32_attention: bool = False, eps: float = 1e-12,
+                        ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    """
+    Compute attention with feature maps by instantiating L x L matrix of attention weights
+    -> Use for attention distillation
+    -> Assume q, k are shape (batch_size, num_heads, seq_len, feature_dim); v is shape (b, h, l, head_dim)
+    """
+    y = None
+    dtype = q.dtype
+    if fp32_attention:
+        q, k = q.float(), k.float()
+    a = torch.einsum('bhmd,bhnd->bhmn', q, k)  # note we don't scale, tho we could
+    if causal:  # Apply causal mask
+        m, n = a.shape[-2:]
+        causal_mask = torch.ones((m, n), device = a.device, dtype = torch.bool).triu(n - m + 1)
+        a = a.masked_fill(causal_mask, 0)
+    # Normalize to compute attention
+    a = a / (a.sum(dim=-1, keepdim=True) + eps)
+    a = a.to(dtype=dtype) if fp32_attention else a
+    if torch.isnan(a).sum() > 0:
+        breakpoint()
+    if v is not None:
+        y = torch.einsum('bhmn,bhnd->bhmd', a, v)
+    return y, a, None
+
+
+# ---------------------
+# Attention layer class
+# ---------------------
+
+class LolcatsLinearAttention(nn.Module):
+    """
+    LoLCATs attention implementation initialized from a 
+    `LlamaAttention` or `MistralAttention` object (base_attn)
+
+    Most of the arguments are directly tied to argparse args
+    - For now we don't support padding.
+    """
+    def __init__(self,
+                 base_attn: nn.Module,  # like LlamaAttention
+                 feature_map: str,
+                 feature_map_kwargs: dict,
+                 layer_idx: Optional[int] = None,
+                 max_layer_idx: Optional[int] = None,
+                 learned_kernel: Optional[str] = None,
+                 learned_kernel_kwargs: Optional[dict] = None,
+                 tie_qk_kernels: Optional[bool] = False,
+                 rotary_config: Optional[dict] = None,
+                 train_attention: Optional[bool] = False,
+                 remove_base_attn: Optional[bool] = True,
+                 attention_type: Optional[str] = 'lolcats_llama',
+                 mask_value: int = 0,
+                 eps: float = 1e-12,
+                 fp32_attention: bool = False,
+                 track_state_grads: bool = False,
+                 rank: Optional[int] = 0,
+                 **kwargs: any) -> None:
+        super().__init__()
+        self.base_config = getattr(base_attn, 'config', None)
+        if self.base_config is not None:
+            self.base_config = self.base_config.to_dict()
+        self.attention_type = attention_type
+        self.mask_value = mask_value
+        self.eps = eps
+        self.layer_idx = (layer_idx if layer_idx is not None else base_attn.layer_idx)
+        self.max_layer_idx = max_layer_idx
+        self.tie_qk_kernels = tie_qk_kernels
+        self.train_attention = train_attention
+        self.base_inference = False
+        self.fp32_attention = fp32_attention
+        self.track_state_grads = track_state_grads
+        if rank == 0:  # multi-gpu
+            if fp32_attention and layer_idx == 0:
+                print(f'-> fp32_attention is {fp32_attention}')
+            if layer_idx == 0 and feature_map_kwargs is not None:
+                for k, v in feature_map_kwargs.items():
+                    print(f'-> {k}: {v}')
+            if layer_idx == 0 and learned_kernel_kwargs is not None:
+                for k, v in learned_kernel_kwargs.items():
+                    print(f'-> {k}: {v}')
+                    
+        self.remove_base_attn = remove_base_attn
+
+        # Rotary embeddings (patch for Llama 3.1, Transformer v4.43.0)
+        self.rotary_config = rotary_config
+        if isinstance(self.rotary_config, DictConfig):  # ensure dict
+            self.rotary_config = OmegaConf.to_container(self.rotary_config)
+        
+        self.rotary_emb = None
+        if self.base_config is not None and self.rotary_config is None:
+            self.rotary_emb = base_attn.rotary_emb
+
+        self.init_weights_(base_attn, remove_base_attn)
+        self.init_feature_map_(feature_map, feature_map_kwargs,
+                               learned_kernel, learned_kernel_kwargs)
+
+    def init_feature_map_(self,
+                          feature_map: str,
+                          feature_map_kwargs: dict,
+                          learned_kernel: str = None,
+                          learned_kernel_kwargs: dict = None):
+        """
+        Initialize MLP-based feature map
+        """
+        self.fmap_gqa = False  # Turn True if specified below
+        if learned_kernel is not None:
+            # Ensure dict
+            learned_kernel_kwargs = {k: v for k, v in learned_kernel_kwargs.items()}
+            learned_kernel_kwargs['num_heads'] = self.num_heads
+            learned_kernel_kwargs['head_dim']  = self.head_dim
+            learned_kernel_kwargs['dtype']     = self.q_proj.weight.dtype
+            learned_kernel_kwargs['device']    = self.q_proj.weight.device
+            # Create MLP
+            mlp_learned_kernel = init_learned_kernel(learned_kernel, **learned_kernel_kwargs)
+        # Add "activation"; see src.models.feature_map.py
+        self.feature_map_q = init_feature_map(name=feature_map,
+                                              mlp=mlp_learned_kernel,
+                                              **feature_map_kwargs)
+        if self.tie_qk_kernels:  # tie mlp weights for query and key feature maps
+            self.feature_map_k = self.feature_map_q
+        else:
+            self.feature_map_k = copy.deepcopy(self.feature_map_q)
+
+    def init_weights_(self, base_attn: nn.Module, remove_base_attn: bool = True):
+        """
+        Initialize module layers, weights, positional dependencies, etc. 
+        from original softmax attention layer (base_attn)
+        """
+        # Make other attributes accessible
+        self.attention_dropout = 0  # We don't use dropout
+        self.hidden_size = base_attn.hidden_size
+        self.num_heads = base_attn.num_heads
+        self.head_dim = base_attn.head_dim
+        self.num_key_value_heads = base_attn.num_key_value_heads
+        self.num_key_value_groups = base_attn.num_key_value_groups
+
+        self.q_shape = [self.num_heads, self.head_dim]
+        self.k_shape = [self.num_key_value_heads, self.head_dim]
+        self.v_shape = [self.num_key_value_heads, self.head_dim]
+        device = base_attn.q_proj.weight.device
+        # Rotary embeddings
+        if self.rotary_emb is None:
+            self.max_position_embeddings = base_attn.max_position_embeddings
+            scaling_factor = getattr(base_attn.rotary_emb, 'scaling_factor', 1.)
+            if self.rotary_config is None:
+                self.rotary_emb = get_rotary_embeddings(
+                    rope_scaling_type=None,
+                    head_dim=self.head_dim,
+                    max_position_embeddings=self.max_position_embeddings,  # base_attn.rotary_emb.max_position_embeddings,
+                    rope_theta=base_attn.rotary_emb.base,
+                    rope_scaling_factor=scaling_factor,  # base_attn.rotary_emb.scaling_factor,
+                    device=device,
+                )
+            else:
+                if 'device' not in self.rotary_config:
+                    self.rotary_config['device'] = device
+                self.rotary_emb = get_rotary_embeddings(**self.rotary_config)
+
+        # Copy original model projection layers
+        self.q_proj = base_attn.q_proj
+        self.k_proj = base_attn.k_proj
+        self.v_proj = base_attn.v_proj
+        self.o_proj = base_attn.o_proj
+        try:  # If wanting to use FA2 for ground-truth inference
+            self._flash_attn_uses_top_left_mask = base_attn._flash_attn_uses_top_left_mask
+        except AttributeError: 
+            pass
+
+        if self.remove_base_attn or remove_base_attn:
+            del base_attn  # We don't need to keep these around
+        else:
+            self.base_attn = base_attn  # For some training runs helpful to just call
+
+    def process_qkv(self,
+                    hidden_states: torch.Tensor,
+                    attention_mask: Optional[torch.Tensor] = None,
+                    position_ids: Optional[torch.LongTensor] = None,
+                    past_key_value: Optional[Tuple[int, torch.Tensor, torch.Tensor]] = None,):  # "legacy" cache approach
+        """
+        Compute queries, keys, and values
+        """
+        b, l, _ = hidden_states.size()
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+        kv_seq_len = k.shape[-2]
+
+        # Shape is (batch_size, seq_len, num_heads, head_dim)
+        q = q.view(b, l, *self.q_shape).transpose(1, 2)
+        k = k.view(b, l, *self.k_shape).transpose(1, 2)
+        v = v.view(b, l, *self.v_shape).transpose(1, 2)
+
+        if past_key_value is not None:  #  and k.shape[2] > q.shape[2]:  # e.g., when generating
+            past_key_value.window_size = getattr(self, 'decode_window_size', None)  # self.decode_window_size
+            if isinstance(past_key_value, Cache):  # In Transformers v4.36+ this is a DynamicCache object
+                kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+            else:
+                kv_seq_len += past_key_value[0].shape[-2]
+
+        # Apply rotary embeddings and repeat for GQA
+        if position_ids is not None and kv_seq_len <= position_ids[0, -1]:
+            kv_seq_len = position_ids[0, -1] + 1  # hack for adjusting position ids
+        try: # As in Transformers v4.36
+            cos, sin = self.rotary_emb(k, seq_len=kv_seq_len)
+            q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
+        except TypeError:  # As in Transformers v4.39+
+            cos, sin = self.rotary_emb(v, position_ids)
+            q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        k = repeat_kv(k, self.num_key_value_groups)
+        v = repeat_kv(v, self.num_key_value_groups)
+        return q, k, v, kv_seq_len
+
+    def forward(self,
+                hidden_states: torch.Tensor,
+                attention_mask: Optional[torch.Tensor] = None,
+                position_ids: Optional[torch.LongTensor] = None,
+                past_key_value: Optional[Tuple[int, torch.Tensor, torch.Tensor]] = None,  # "legacy" cache approach
+                output_attentions: bool = False,
+                use_cache: bool = False,
+                **kwargs,
+               ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """
+        Forward pass modified from transformers.models.mistral.modeling_mistral (v4.36)
+        - Consistent with HuggingFace Transformers for easy use with their pretrained models
+        """
+        b, l, _ = hidden_states.size()
+        q, k, v, kv_seq_len = self.process_qkv(hidden_states, attention_mask, 
+                                               position_ids, past_key_value)
+        if self.base_inference:
+            with torch.no_grad():
+                # 1. Compute "ground-truth" attention output and weights
+                y_true, _, _ = softmax_attention(q, k, v, causal=True)
+                y_true = y_true.transpose(1, 2).contiguous().view(b, l, self.hidden_size)
+                y_true = self.o_proj(y_true)
+                attn_weights = (None, None)
+
+        elif self.train_attention:  # Distilling / learning attentions
+            # Note for now we assume no padding when distilling; attention masks only enforce causality
+            assert output_attentions is True, f'When training feature maps, output_attentions should be True but is {output_attentions}'
+            with torch.no_grad():
+                # 1. Compute "ground-truth" attention output and weights
+                _y_true, attn_true, _ = softmax_attention(q, k, v, causal=True)
+                y_true = _y_true.transpose(1, 2).contiguous().view(b, l, self.hidden_size)
+                y_true = self.o_proj(y_true)
+        
+            # 2. Compute "predicted" attention (just weights)
+            q, k = self.feature_map_q.q_map(q), self.feature_map_k.k_map(k)
+            y_pred, attn_pred, _ = quadratic_attention(q, k, v, causal=True)
+            attn_weights = ((attn_pred, attn_true), (y_pred, _y_true))  # Save both attention weights so we can supervise. 
+        
+        else:  # Finetuning
+            q, k = self.feature_map_q(q), self.feature_map_k(k)
+            # Apply prefill mask
+            if attention_mask is not None and q.shape[2] > 1:
+                if len(attention_mask.shape) == 4:
+                    lin_attn_mask = (attention_mask == 0)[:, :1, -1, :l][..., None]  # b, 1, k_len, 1
+                else:
+                    lin_attn_mask = attention_mask[:, None, :, None]  # b, 1, k_len, 1
+                k = k.masked_fill(~lin_attn_mask, 0)
+            
+            if past_key_value is not None:  # Initialize states
+                if len(past_key_value.kv_states) == self.layer_idx:
+                    b, h, _, f = k.shape
+                    past_key_value.kv_states.append(
+                        torch.zeros(b, h, f, self.head_dim, dtype=q.dtype, device=q.device)
+                    )
+                    past_key_value.k_states.append(
+                        torch.zeros(b, h, 1, f, dtype=q.dtype, device=q.device)
+                    )
+                # Generating
+                if q.shape[2] == 1 and kv_seq_len > 1 and past_key_value is not None:
+                    assert use_cache is True
+                    kv_state, k_state = past_key_value.update(k, v, self.layer_idx,
+                                                              accumulate_in_fp32=self.fp32_attention)
+                    if self.fp32_attention:
+                        q = q.float()
+                        y_true = (torch.einsum('bhlf,bhfd->bhld', q, kv_state.float()) /
+                                  torch.einsum('bhlf,bhlf->bhl', q, k_state.float())[..., None]).to(dtype=k.dtype)
+                    else:
+                        y_true = (torch.einsum('bhlf,bhfd->bhld', q, kv_state) /
+                                  torch.einsum('bhlf,bhlf->bhl', q, k_state)[..., None])
+                else:
+                    kv_state = past_key_value.kv_states[self.layer_idx]
+                    k_state  = past_key_value.k_states[self.layer_idx]
+                    y_true, _, _ = linear_attention(q, k, v, self.fp32_attention, self.eps)  # Ordinarily the states are ignored
+                    past_key_value.update(k.detach(), v.detach(), self.layer_idx,
+                                          accumulate_in_fp32=self.fp32_attention)
+                                          # doing some unnecessary recomputation here
+            else:
+                y_true, _, _ = linear_attention(q, k, v, self.fp32_attention, self.eps)
+
+            # Concatenate heads and apply output projection
+            y_true = y_true.transpose(1, 2).contiguous().view(b, l, self.hidden_size)
+            y_true = self.o_proj(y_true)
+            attn_weights = None
+
+        return y_true, attn_weights, past_key_value
+
+
+class LinearAttentionState(Cache):
+    """
+    Handle the KV and K states for linear attention
+    - Adopts HF Transformers `past_key_values` convention
+    - Inherits from `Cache` class
+    - Modified from transformers.cache_utils.DynamicCache (v4.36)
+    """
+    def __init__(self) -> None:
+        self._seen_tokens = 0  # should be `self.seen_tokens` in Transformers v4.36
+        self._seen_tokens_by_layer: List[int] = []
+        self.kv_states: List[torch.Tensor] = []
+        self.k_states:  List[torch.Tensor] = []
+
+    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
+        """
+        Returns the sequence length of the cached states. A layer index can be optionally passed.
+        """
+        if len(self._seen_tokens_by_layer) <= layer_idx:  # Initializing kv and k states
+            self._seen_tokens_by_layer.append(0)
+        return self._seen_tokens_by_layer[layer_idx]
+
+    def get_max_length(self) -> Optional[int]:
+        """
+        Returns the maximum sequence length of the cached states. DynamicCache does not have a maximum length.
+        """
+        return None
+
+    def get_usable_length(self, new_seq_length: int, layer_idx: Optional[int] = 0) -> int:
+        """Given the sequence length of the new inputs, returns the usable length of the cache."""
+        # Cache without size limit -> all cache is usable
+        # Cache with size limit -> if the length cache plus the length of the new inputs is larger the maximum cache
+        #   length, we will need to evict part of the cache (and thus not all cache is usable)
+        max_length = self.get_max_length()
+        previous_seq_length = self.get_seq_length(layer_idx)
+        if max_length is not None and previous_seq_length + new_seq_length > max_length:
+            return max_length - new_seq_length
+        return previous_seq_length
+
+    def update(self, key_states: torch.Tensor, value_states: torch.Tensor,
+               layer_idx: Optional[int] = None, cache_kwargs: Optional[any] = None,
+               accumulate_in_fp32: bool = True, **kwargs: any,
+              ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        with torch.no_grad ():
+            if layer_idx == 0:
+                self._seen_tokens += key_states.shape[-2]
+            dtype = key_states.dtype
+            if accumulate_in_fp32:
+                key_states, value_states = key_states.float(), value_states.float()
+
+            kv_state = torch.einsum('bhlf,bhld->bhfd', key_states, value_states).detach()
+            k_state  = key_states.sum(dim=-2, keepdim=True).detach()  # b, h, 1, f; note the 1
+            # Update the cache
+            if len(self.k_states) <= layer_idx:  # Initializing kv and k states
+                print('if len(self.k_states) <= layer_idx:  # Initializing kv and k states')
+                self.kv_states.append(kv_state.to(dtype))
+                self.k_states.append(k_state.to(dtype))
+            else:
+                kv_state = (self.kv_states[layer_idx].to(kv_state.dtype) + kv_state).to(dtype)
+                k_state  = (self.k_states[layer_idx].to(kv_state.dtype) + k_state).to(dtype)
+                self.kv_states[layer_idx] = kv_state
+                self.k_states[layer_idx]  = k_state
+            self._seen_tokens_by_layer[layer_idx] += key_states.shape[-2] 
+        return self.kv_states[layer_idx], self.k_states[layer_idx]
+
+    def to_legacy_cache(self):
+        """Hack, but just return self"""
+        return self
+
+    def reorder_cache(self, beam_idx: torch.LongTensor):
+        """
+        Reorders the cache for beam search, given the selected beam indices.
+        -> Copied from transformers/src/transformers/cache_utils.py
+        """
+        raise NotImplementedError('Reordering cache not implemented for LinearAttentionState')

src/model/linear_attention/linear_window_attention_sw.py

@@ -0,0 +1,339 @@
+"""
+Subquadratic attention combining sliding window and linear attentions
+- Using "standard" sliding windows
+- Didactically computes outputs with n^2 attention weights for now
+- Copied + adapted from linear_window_attention_tk.py for single-file reference
+
+For each layer: 
+- We first compute (softmax) attention over sliding windows
+- We then compute standard linear attention to "fill in" the earlier parts
+- We combine to model the entire sequence
+"""
+from typing import List, Tuple, Optional, Callable
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.cache_utils import Cache
+
+from .linear_attention import (
+    LolcatsLinearAttention, LinearAttentionState, 
+    softmax_attention
+)
+
+# ----------------------
+# Sliding window helpers
+# ----------------------
+def get_masks(window_size: int, q_len: int, k_len: int, 
+              device: torch.device) -> tuple[torch.Tensor]:
+    """
+    Return masks for softmax and linear attention terms
+    -> 1 is include, 0 is ignore
+    """
+    kwargs = {'device': device, 'dtype': int}
+    causal_mask = torch.ones((q_len, k_len), **kwargs).tril(k_len - q_len)
+    linear_mask = torch.ones((q_len, k_len), **kwargs).tril(k_len - q_len - window_size)
+    window_mask = causal_mask - linear_mask
+    # Return softmax mask (window), linear attention mask
+    # -> shapes broadcast over (b, h, q_len, k_len)
+    return window_mask[None, None, ...], linear_mask[None, None, ...]
+
+
+def hybrid_attention_quadratic(q: torch.Tensor, k: torch.Tensor, 
+                               f_q: torch.Tensor, f_k: torch.Tensor,
+                               v: torch.Tensor,
+                               window_factor: torch.Tensor,
+                               linear_factor: torch.Tensor,
+                               window_size: int,
+                               kv_state: torch.Tensor = None,
+                               k_state: torch.Tensor = None,
+                               eps: float = 1e-12,
+                               mask_value: float=-1e8):
+    """
+    Hybrid attention combining sliding window and linear attentions
+    """
+
+    mask_window, mask_linear = get_masks(window_size, q.shape[-2], k.shape[-2], q.device)
+
+    # 1. Sliding window (softmax attention)
+    a_sm = torch.einsum('bhmd,bhnd->bhmn', q.float(), k.float()) * (k.shape[-1] ** -0.5)
+    a_sm = a_sm.masked_fill(~mask_window.bool(), mask_value)
+    # torch.softmax(a_sm, dim=-1), but we account for the max when combining
+    a_sm_max = torch.amax(a_sm, dim=-1, keepdim=True)
+    a_sm   = window_factor * torch.exp(a_sm - a_sm_max)
+    sum_sm = a_sm.sum(dim=-1, keepdim=True)
+
+    # 2. Under window (linear attention)
+    a_ln = torch.einsum('bhmd,bhnd->bhmn', f_q.float(), f_k.float())
+    a_ln = linear_factor * a_ln.masked_fill(~mask_linear.bool(), 0)
+    sum_ln = a_ln.sum(dim=-1, keepdim=True)
+
+    # 3. Combine
+    a = ((a_sm + a_ln) / (sum_sm + sum_ln)).to(q.dtype)  # Save attention weights
+    # Allow outputs to also depend on prior kv_state and k_state
+    y = torch.einsum('bhmn,bhnd->bhmd', a_sm + a_ln, v.float())
+    if kv_state is not None:  # Combine with prior kv_state and k_state
+        y += linear_factor * torch.einsum('bhld,bhdf->bhlf', f_q.float(), kv_state.float())
+        sum_ln += linear_factor * torch.einsum(
+            'bhld,bhnd->bhl', f_q.float(), k_state.float())[..., None]
+    y = (y / (sum_sm + sum_ln)).to(q.dtype)
+    return y, a  # attention weights only for the last chunk
+
+
+# ---------------------
+# Attention layer class
+# ---------------------
+class LolcatsSlidingWindowAttention(LolcatsLinearAttention):
+    """
+    Lolcats attention combining sliding window and linear attention
+    """
+    def __init__(self, 
+                 window_size: int = 64, 
+                 decode_window_size: int = None,
+                 affine_attention_factors: bool = False,
+                 init_window_factor: float = 0,
+                 train_window_factor: bool = True,
+                 state_grad_enabled: bool = False,
+                 **kwargs):
+        self.window_size = window_size
+        self.decode_window_size = (
+            decode_window_size if decode_window_size is not None else window_size
+        )
+        self.window_kwargs = {'dimension': 2, 'size': window_size, 'step': 1}
+        super().__init__(**kwargs)
+        self.attention_type = kwargs['attention_type']  #  'hedgehog_llama_window_sw'
+        # Determine how we compute attentions
+        self.quadratic_attention = hybrid_attention_quadratic
+        self.attention_type = kwargs['attention_type']  # 'hedgehog_long_llama_window_sw'
+        # Learnable factor for combining attentions
+        self.affine_attention_factors = affine_attention_factors
+        device, dtype = self.q_proj.weight.device, self.q_proj.weight.dtype
+        if train_window_factor:
+            self.window_factors = nn.Parameter(
+                init_window_factor * torch.ones(1, self.num_heads, 1, 1, device=device, dtype=dtype))
+        else:
+            self.register_buffer(
+                "window_factors", init_window_factor * torch.ones(1, self.num_heads, 1, 1, device=device, dtype=dtype)
+            )
+        # Whether we use original flash attention 2 inference (use during attention transfer)
+        self.base_inference = False
+        self.state_grad_enabled = state_grad_enabled
+        
+    def forward(self,
+                hidden_states: torch.Tensor,
+                attention_mask: Optional[torch.Tensor] = None,
+                position_ids: Optional[torch.LongTensor] = None,
+                past_key_value: Optional[Cache] = None,
+                output_attentions: bool = False,
+                use_cache: bool = False,
+                **kwargs,
+               ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """
+        Forward pass with the option to compute attention weights multiple ways
+        if self.train_attention is True
+        -> Consistent with HuggingFace Transformers for easy use with their pretrained models
+        """
+        b, l, _ = hidden_states.size()
+        q, k, v, kv_seq_len = self.process_qkv(hidden_states, attention_mask, 
+                                               position_ids, past_key_value)
+        f_q, f_k = self.feature_map_q(q), self.feature_map_k(k)  # Have to do after repeat for grouped-query attn if we use same fmap
+
+        if self.train_attention:
+            # 1. Compute "ground-truth" attention output and weights
+            with torch.no_grad():
+                _y_true, a_true = softmax_attention(q, k, v)[:2]
+                y_true = _y_true.transpose(1, 2).contiguous().view(b, l, self.hidden_size)
+                y_true = self.o_proj(y_true)
+
+            # 2. Compute "predicted" attention outputs
+            # compute attn weights under sliding window
+            window_factors = F.sigmoid(self.window_factors)
+            linear_factors = 1 - window_factors if self.affine_attention_factors else 1
+            y_pred, a_pred = self.quadratic_attention(q, k, f_q, f_k, v,
+                                                      window_factors, linear_factors,
+                                                      window_size=self.window_size)
+            attn_weights = ((a_pred, a_true), (y_pred, _y_true))
+        else:
+            attn_weights = None
+            # attention_mask = None  # For now this is always True
+            if past_key_value is None:  # Regular training
+                window_factors = F.sigmoid(self.window_factors)
+                linear_factors = 1 - window_factors if self.affine_attention_factors else 1
+                y_true, a_pred = self.quadratic_attention(q, k, f_q, f_k, v,
+                                                          window_factors, linear_factors,
+                                                          window_size=self.window_size)
+                attn_weights = a_pred
+            else:
+                past_key_value.window_size = self.decode_window_size
+                if f_q.shape[2] == 1 and kv_seq_len > 1 and not self.training:  # Generating
+                    assert use_cache is True
+                    _kv = past_key_value.update_for_decoding(k, v, self.layer_idx,
+                                                             self.feature_map_k,
+                                                             dtype=q.dtype)
+                    k_cache, v_cache, f_kv_state, f_k_state = _kv
+
+                    # Sliding window + linear attention decode
+                    window_factors = F.sigmoid(self.window_factors)
+                    linear_factors = 1 - window_factors if self.affine_attention_factors else 1
+
+                    # Softmax attention terms
+                    a_sm = torch.einsum('bhmd,bhnd->bhmn', q.float(), k_cache.float()) * (k.shape[-1] ** -0.5)
+                    a_sm_max = torch.amax(a_sm, dim=-1, keepdim=True)
+                    a_sm   = window_factors * torch.exp(a_sm - a_sm_max)
+                    sum_sm = a_sm.sum(dim=-1, keepdim=True)
+
+                    # Combine with linear attention terms
+                    y_true = (torch.einsum('bhmn,bhnd->bhmd', a_sm, v_cache.float())
+                              + linear_factors * torch.einsum('bhlf,bhfd->bhld', f_q.float(), f_kv_state.float()))
+                    sum_ln = linear_factors * torch.einsum(
+                        'bhlf,bhnf->bhl', f_q.float(), f_k_state.float())[..., None]
+                    y_true = (y_true / (sum_sm + sum_ln)).to(q.dtype) 
+
+                else:  # Stateful training
+                    try:
+                        kv_state = past_key_value.kv_states[self.layer_idx]
+                        k_state  = past_key_value.k_states[self.layer_idx]
+                    except IndexError:
+                        kv_state, k_state = None, None
+                    window_factors = F.sigmoid(self.window_factors)
+                    linear_factors = 1 - window_factors if self.affine_attention_factors else 1
+                    y_true, _ = self.quadratic_attention(q, k, f_q, f_k, v,
+                                                         window_factors, linear_factors,
+                                                         window_size=self.window_size,
+                                                         kv_state=kv_state,
+                                                         k_state=k_state)
+                    # Save and update KV cache and states
+                    # past_key_value.update(k, v.detach(), self.layer_idx,
+                    #                       fmap_key_states=f_k.detach(),
+                    #                       accumulate_in_fp32=True)
+                    past_key_value.update(k, v, self.layer_idx,
+                                          fmap_key_states=f_k,
+                                          accumulate_in_fp32=True)
+            # Concatenate heads and apply output projection
+            y_true = y_true.transpose(1, 2).contiguous().view(b, l, self.hidden_size)
+            y_true = self.o_proj(y_true)
+        return y_true, attn_weights, past_key_value
+
+
+class LinearAttentionSlidingWindowCache(LinearAttentionState):
+    """
+    Class for `past_key_values`
+    -> Alternative to KV cache; here we only maintain a "KV state" and "K state"
+    -> Modified from transformers.cache_utils.DynamicCache (v4.36)
+    """
+    def __init__(self, window_size: int = 64) -> None:
+        super().__init__()
+        self._seen_tokens = 0  # should be `self.seen_tokens` in Transformers v4.36
+        self._seen_tokens_by_layer: List[int] = []
+        self.kv_states: List[torch.Tensor] = []
+        self.k_states:  List[torch.Tensor] = []
+
+        # Account for sliding windows
+        self.decode_kv_states: List[torch.Tensor] = []
+        self.decode_k_states: List[torch.Tensor] = []
+        self.k_cache: List[torch.Tensor] = []
+        self.v_cache: List[torch.Tensor] = []
+        self.window_size = window_size
+
+    def update(self, key_states: torch.Tensor, value_states: torch.Tensor, 
+               layer_idx: Optional[int] = None, cache_kwargs: Optional[any] = None,
+               accumulate_in_fp32: bool = False, 
+               fmap_key_states: torch.Tensor = None,  # should not be None
+               grad_enabled: bool = False,
+               **kwargs: any,
+              ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Update KV, K states; and KV cache during training
+        - For decoding, use `self.decode_kv_states` to keep track of KV states 
+          up to sliding window terms
+        - For (chunked) training, use `self.kv_states` to keep track of KV states
+          up to end of sequence
+        - Likewise for `self.decode_k_states` and `self.k_states`
+        """
+        with torch.set_grad_enabled(grad_enabled):
+            if layer_idx == 0:
+                self._seen_tokens += key_states.shape[-2]
+
+            dtype = key_states.dtype
+            if accumulate_in_fp32:
+                # key_states = key_states.float()
+                fmap_key_states = fmap_key_states.float()
+                value_states = value_states.float()
+
+            # Decoding KV state (KV terms up to last window_size)
+            decode_kv_state = torch.einsum(
+                'bhlf,bhld->bhfd', fmap_key_states[:, :, :-self.window_size], value_states[:, :, :-self.window_size]
+            )
+            # KV state
+            kv_state = decode_kv_state + torch.einsum(
+                'bhlf,bhld->bhfd', fmap_key_states[:, :, -self.window_size:], value_states[:, :, -self.window_size:]
+            )
+            # shape is b, h, 1, f; note the 1
+            decode_k_state = fmap_key_states[:, :, :-self.window_size].sum(dim=-2, keepdim=True)
+            k_state = (decode_k_state + fmap_key_states[:, :, -self.window_size:].sum(dim=-2, keepdim=True))
+
+            # Update the cache
+            if len(self.k_states) <= layer_idx:  # Initializing kv and k states
+                self.kv_states.append(kv_state.to(dtype))
+                self.k_states.append(k_state.to(dtype))
+
+                self.decode_kv_states.append(decode_kv_state.to(dtype))
+                self.decode_k_states.append(decode_k_state.to(dtype))
+
+                self.k_cache.append(key_states[:, :, -self.window_size:, :])
+                self.v_cache.append(value_states[:, :, -self.window_size:, :].to(dtype))
+                # self._seen_tokens_by_layer[layer_idx].append(key_states.shape[-2])
+            else:
+                # Update kv and k states recurrently
+                kv_state = (self.kv_states[layer_idx].to(kv_state.dtype) + kv_state).to(dtype)
+                k_state  = (self.k_states[layer_idx].to(kv_state.dtype) + k_state).to(dtype)
+                self.kv_states[layer_idx] = kv_state
+                self.k_states[layer_idx]  = k_state
+
+                decode_kv_state = (self.decode_kv_states[layer_idx].to(kv_state.dtype) 
+                                   + decode_kv_state).to(dtype)
+                decode_k_state  = (self.decode_k_states[layer_idx].to(kv_state.dtype) 
+                                   + decode_k_state).to(dtype)
+                self.decode_kv_states[layer_idx] = decode_kv_state
+                self.decode_k_states[layer_idx]  = decode_k_state
+
+                self.k_cache[layer_idx] = key_states[:, :, -self.window_size:, :]
+                self.v_cache[layer_idx] = value_states[:, :, -self.window_size:, :]
+            self._seen_tokens_by_layer[layer_idx] += key_states.shape[-2]
+
+        return self.kv_states[layer_idx], self.k_states[layer_idx]
+
+    def update_for_decoding(self, keys: torch.Tensor, values: torch.Tensor, 
+                            layer_idx: int, feature_map_k: Callable, dtype: torch.dtype):
+        """
+        Update the decoding KV and K states, and KV cache, during decodeing
+        """
+        with torch.no_grad():
+            k_cache = self.k_cache[layer_idx]
+            v_cache = self.v_cache[layer_idx]
+
+            if k_cache.shape[-2] < self.window_size:  # build window-size cache
+                self.k_cache[layer_idx] = torch.cat([k_cache, keys], dim=-2)
+                self.v_cache[layer_idx] = torch.cat([v_cache, values], dim=-2)
+            else:
+                # MZ 6/3: handle short inputs; zero-out padding when initial k.shape[2] < self.window_size
+                # if k_cache[:, :, :1, :].sum() == 0:   # heuristic for zeroing out padding in cache
+                #     f_k_state = torch.zeros(k_cache[:, :, :1, :].shape, dtype=dtype, device=k_cache.device)
+                # else:
+                #     f_k_state = feature_map_k(k_cache[:, :, :1, :])
+                # -> MZ (later): above only relevant if we zero-pad in our hybrid attention computation
+                k_state = feature_map_k(k_cache[:, :, :1, :])
+                v_state = v_cache[:, :, :1, :]
+                kv_state = torch.einsum('bhlf,bhld->bhfd', k_state.float(), v_state.float()).to(dtype) # b, h, f, d
+                self.decode_kv_states[layer_idx] += kv_state
+                self.decode_k_states[layer_idx] += k_state
+                
+                self.k_cache[layer_idx] = torch.cat([k_cache[:, :, 1:, :], keys], dim=-2)
+                self.v_cache[layer_idx] = torch.cat([v_cache[:, :, 1:, :], values], dim=-2)
+            
+            if layer_idx == 0:
+                self._seen_tokens += keys.shape[-2]
+            self._seen_tokens_by_layer[layer_idx] += keys.shape[-2]
+            return (self.k_cache[layer_idx], self.v_cache[layer_idx], 
+                    self.decode_kv_states[layer_idx], self.decode_k_states[layer_idx])

src/model/linear_attention/linear_window_attention_sw_linear.py

@@ -0,0 +1,522 @@
+"""
+Subquadratic attention combining sliding window and linear attentions
+- Using "standard" sliding windows
+- Didactically computes outputs with n^2 attention weights for now
+- Copied + adapted from linear_window_attention_tk.py for single-file reference
+
+For each layer: 
+- We first compute (softmax) attention over sliding windows
+- We then compute standard linear attention to "fill in" the earlier parts
+- We combine to model the entire sequence
+"""
+from typing import List, Tuple, Optional, Callable
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.cache_utils import Cache
+try:
+    from transformers.modeling_flash_attention_utils import _flash_attention_forward
+except ModuleNotFoundError:
+    _flash_attention_forward = None  # Transformers v4.36
+
+# Causal linear attention dot product CUDA kernel from fast-transformers
+from csrc import causal_dot_product
+
+from src.model.rotary import apply_rotary_pos_emb
+from .linear_attention import (
+    LolcatsLinearAttention, LinearAttentionState, 
+    softmax_attention
+)
+
+# ----------------------
+# Sliding window helpers
+# ----------------------
+def get_masks(window_size: int, q_len: int, k_len: int, 
+              device: torch.device) -> tuple[torch.Tensor]:
+    """
+    Return masks for softmax and linear attention terms
+    -> 1 is include, 0 is ignore
+    """
+    kwargs = {'device': device, 'dtype': int}
+    causal_mask = torch.ones((q_len, k_len), **kwargs).tril(max(k_len - q_len, 0))
+    linear_mask = torch.ones((q_len, k_len), **kwargs).tril(max(k_len - q_len, 0) - window_size)
+    window_mask = causal_mask - linear_mask
+    # Return softmax mask (window), linear attention mask
+    # -> shapes broadcast over (b, h, q_len, k_len)
+    return window_mask[None, None, ...], linear_mask[None, None, ...]
+
+
+def hybrid_attention_quadratic(q: torch.Tensor, k: torch.Tensor, 
+                               f_q: torch.Tensor, f_k: torch.Tensor,
+                               v: torch.Tensor,
+                               window_factor: torch.Tensor,
+                               linear_factor: torch.Tensor,
+                               window_size: int,
+                               kv_state: torch.Tensor = None,
+                               k_state: torch.Tensor = None,
+                               eps: float = 1e-12,
+                               mask_value: float=-1e8):
+    """
+    Hybrid attention combining sliding window and linear attentions
+    """
+
+    mask_window, mask_linear = get_masks(window_size, q.shape[-2], k.shape[-2], q.device)
+
+    # 1. Sliding window (softmax attention)
+    a_sm = torch.einsum('bhmd,bhnd->bhmn', q.float(), k.float()) * (k.shape[-1] ** -0.5)
+    a_sm = a_sm.masked_fill(~mask_window.bool(), mask_value)
+    # torch.softmax(a_sm, dim=-1), but we account for the max when combining
+    a_sm_max = torch.amax(a_sm, dim=-1, keepdim=True)
+    a_sm   = window_factor * torch.exp(a_sm - a_sm_max)
+    sum_sm = a_sm.sum(dim=-1, keepdim=True)
+
+    # 2. Under window (linear attention)
+    a_ln = torch.einsum('bhmd,bhnd->bhmn', f_q.float(), f_k.float())
+    a_ln = linear_factor * a_ln.masked_fill(~mask_linear.bool(), 0)
+    sum_ln = a_ln.sum(dim=-1, keepdim=True)
+
+    # 3. Combine
+    a = ((a_sm + a_ln) / (sum_sm + sum_ln)).to(q.dtype)  # Save attention weights
+    # Allow outputs to also depend on prior kv_state and k_state
+    y = torch.einsum('bhmn,bhnd->bhmd', a_sm + a_ln, v.float())
+    if kv_state is not None:  # Combine with prior kv_state and k_state
+        y += linear_factor * torch.einsum('bhld,bhdf->bhlf', f_q.float(), kv_state.float())
+        sum_ln += linear_factor * torch.einsum(
+            'bhld,bhnd->bhl', f_q.float(), k_state.float())[..., None]
+    y = (y / (sum_sm + sum_ln)).to(q.dtype)
+    return y, a  # attention weights only for the last chunk
+
+
+# ------------------------------
+# Hybrid window attention linear
+# ------------------------------
+def under_window_linear_attention(f_q: torch.Tensor, f_k: torch.Tensor, v: torch.Tensor,
+                                  window_size: int, linear_factor: float, eps: float=1e-12):
+    """Compute hybrid window attention dot product with linear complexity in q_len"""
+    dtype = f_q.dtype
+    w   = window_size
+    f_k = F.pad(f_k, (0, 0, w, 0), value=0)[:, :, :-w, :]
+    v   = F.pad(v, (0, 0, w, 0), value=0)[:, :, :-w, :]
+    qkv = linear_factor * causal_dot_product(f_q.contiguous().to(dtype=torch.float32), 
+                                             f_k.contiguous().to(dtype=torch.float32), 
+                                             v.contiguous().to(dtype=torch.float32)).to(dtype=dtype)
+    sum_f_k = f_k.float().cumsum(dim=2).to(dtype=dtype)
+    sum_qk = linear_factor * torch.einsum("bhld,bhld->bhl", f_q, sum_f_k)[..., None]
+    sum_qk[sum_qk == 0] += eps
+    return qkv, sum_qk
+
+
+def sliding_window_softmax_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+                                     window_size: int, window_factor: float, mask_value: float=-1e8):
+    """
+    Compute sliding window softmax attention without materializing 
+    O(seq_len^2) attention weights
+    """        
+    d = q.shape[-1]
+    # Compute windows for keys
+    window_kwargs = {'dimension': 2, 'size': window_size, 'step': 1}
+    k = F.pad(k, (0, 0, window_size - 1, 0), value=0).unfold(**window_kwargs)
+    v = F.pad(v, (0, 0, window_size - 1, 0), value=0).unfold(**window_kwargs)
+
+    # Compute windowed_softmax(qk); causal in its construction
+    a_sm = torch.einsum('bhld,bhldw->bhlw', q, k) * (d ** -0.5)
+    a_sm[a_sm == 0] = -torch.finfo(q.dtype).max  # heuristic for zeroing out padding above
+    a_sm_max = torch.amax(a_sm, dim=-1, keepdim=True)
+    a_sm   = window_factor * torch.exp(a_sm - a_sm_max)
+    sum_sm = a_sm.sum(dim=-1, keepdim=True)
+    return torch.einsum('bhlw,bhldw->bhld', a_sm, v), sum_sm
+    # return torch.einsum('bhlw,bhldw->bhld', torch.softmax(qk, dim=-1), v)
+
+
+def hybrid_attention_linear(q: torch.Tensor, k: torch.Tensor,  
+                            f_q: torch.Tensor, f_k: torch.Tensor, 
+                            v: torch.Tensor,
+                            window_factor: torch.Tensor = None,
+                            linear_factor: torch.Tensor = None,
+                            window_size: int = 64,
+                            kv_state: torch.Tensor = None,
+                            k_state: torch.Tensor = None,
+                            eps: float = 1e-12,
+                            mask_value: float=-1e8):
+    """
+    Alternative hybrid attention combining sliding window and linear attentions
+    -> Uses O(n) memory if n is sequence length by padding and unfolding windows
+    """
+    window_kwargs = {'dimension': 2, 'size': window_size, 'step': 1}
+    # 1. Sliding window (softmax attention)
+    with torch.no_grad():
+        qkv_sm, sum_qk_sm = sliding_window_softmax_attention(q, k, v, window_size, window_factor, mask_value)
+
+    # 2. Under window (linear attention)
+    qkv_ln, sum_qk_ln = under_window_linear_attention(f_q, f_k, v, window_size, linear_factor, eps)
+
+    # 3. Combine
+    y = (qkv_sm + qkv_ln) / (sum_qk_sm + sum_qk_ln)
+    return y, None
+
+
+# ---------------------
+# Attention layer class
+# ---------------------
+class LolcatsLinearSlidingWindowAttention(LolcatsLinearAttention):
+    """
+    Lolcats attention combining sliding window and linear attention
+    """
+    def __init__(self, 
+                 window_size: int = 64, 
+                 decode_window_size: int = None,
+                 affine_attention_factors: bool = False,
+                 init_window_factor: float = 0,
+                 train_window_factor: bool = True,
+                 state_grad_enabled: bool = False,
+                 **kwargs):
+        self.window_size = window_size
+        self.decode_window_size = (
+            decode_window_size if decode_window_size is not None else window_size
+        )
+        self.window_kwargs = {'dimension': 2, 'size': window_size, 'step': 1}
+        super().__init__(**kwargs)
+        # Determine how we compute attentions
+        self.linear_attention = hybrid_attention_linear
+        self.attention_type = 'lolcats_llama_window_sw'
+        # Learnable factor for combining attentions
+        self.affine_attention_factors = affine_attention_factors
+        device, dtype = self.q_proj.weight.device, self.q_proj.weight.dtype
+        if train_window_factor:
+            self.window_factors = nn.Parameter(
+                init_window_factor * torch.ones(1, self.num_heads, 1, 1, device=device, dtype=dtype))
+        else:
+            self.register_buffer(
+                "window_factors", init_window_factor * torch.ones(1, self.num_heads, 1, 1, device=device, dtype=dtype)
+            )
+        # Whether we use original flash attention 2 inference (use during attention transfer)
+        self.base_inference = False
+        self.state_grad_enabled = state_grad_enabled
+        
+    def forward(self,
+                hidden_states: torch.Tensor,
+                attention_mask: Optional[torch.Tensor] = None,
+                position_ids: Optional[torch.LongTensor] = None,
+                past_key_value: Optional[Cache] = None,
+                output_attentions: bool = False,
+                use_cache: bool = False,
+                **kwargs,
+               ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        """
+        Forward pass with the option to compute attention weights multiple ways
+        if self.train_attention is True
+        -> Consistent with HuggingFace Transformers for easy use with their pretrained models
+        """
+        b, l, _ = hidden_states.size()
+
+        if self.train_attention and self.base_inference:
+            with torch.no_grad():
+                _y_true = flash_attention_2(self,  # self.base_attn,
+                                            hidden_states=hidden_states,
+                                            attention_mask=None,
+                                            position_ids=position_ids,
+                                            past_key_value=None,
+                                            output_attentions=False,
+                                            use_cache=False)[0]
+                # _y_true.shape is (batch_size, seq_len, num_heads, head_dim)
+                y_true = _y_true.reshape(b, l, -1).contiguous()
+                y_true = self.o_proj(y_true)
+                # layer_io = (hidden_states, _y_true)  # hack
+                layer_io = (hidden_states.cpu(), _y_true.cpu())  # hack
+                return y_true, layer_io, None
+
+        else:
+            q, k, v, kv_seq_len = self.process_qkv(hidden_states, attention_mask, 
+                                                   position_ids, past_key_value)
+            f_q, f_k = self.feature_map_q(q), self.feature_map_k(k)  # Have to do after repeat for grouped-query attn if we use same fmap
+
+            attn_weights = None
+            # attention_mask = None  # For now this is always True
+            if past_key_value is None:  # Regular training
+                window_factors = F.sigmoid(self.window_factors)
+                linear_factors = 1 - window_factors if self.affine_attention_factors else 1
+                y_true, a_pred = self.linear_attention(q, k, f_q, f_k, v,
+                                                       window_factors, linear_factors,
+                                                       window_size=self.window_size)
+                attn_weights = a_pred
+            else:
+                past_key_value.window_size = self.decode_window_size
+                if f_q.shape[2] == 1 and kv_seq_len > 1 and not self.training:  # Generating
+                    assert use_cache is True
+                    _kv = past_key_value.update_for_decoding(k, v, self.layer_idx,
+                                                             self.feature_map_k,
+                                                             dtype=q.dtype)
+                    k_cache, v_cache, f_kv_state, f_k_state = _kv
+
+                    # Sliding window + linear attention decode
+                    window_factors = F.sigmoid(self.window_factors)
+                    linear_factors = 1 - window_factors if self.affine_attention_factors else 1
+
+                    # Softmax attention terms
+                    a_sm = torch.einsum('bhmd,bhnd->bhmn', q.float(), k_cache.float()) * (k.shape[-1] ** -0.5)
+                    a_sm_max = torch.amax(a_sm, dim=-1, keepdim=True)
+                    a_sm   = window_factors * torch.exp(a_sm - a_sm_max)
+                    sum_sm = a_sm.sum(dim=-1, keepdim=True)
+
+                    # Combine with linear attention terms
+                    y_true = (torch.einsum('bhmn,bhnd->bhmd', a_sm, v_cache.float())
+                              + linear_factors * torch.einsum('bhlf,bhfd->bhld', f_q.float(), f_kv_state.float()))
+                    sum_ln = linear_factors * torch.einsum(
+                        'bhlf,bhnf->bhl', f_q.float(), f_k_state.float())[..., None]
+                    y_true = (y_true / (sum_sm + sum_ln)).to(q.dtype) 
+
+                else:  # Stateful training
+                    try:
+                        kv_state = past_key_value.kv_states[self.layer_idx]
+                        k_state  = past_key_value.k_states[self.layer_idx]
+                    except IndexError:
+                        kv_state, k_state = None, None
+                    window_factors = F.sigmoid(self.window_factors)
+                    linear_factors = 1 - window_factors if self.affine_attention_factors else 1
+                    y_true, _ = self.linear_attention(q, k, f_q, f_k, v,
+                                                      window_factors, linear_factors,
+                                                      window_size=self.window_size,
+                                                      kv_state=kv_state,
+                                                      k_state=k_state)
+                    # Save and update KV cache and states
+                    # past_key_value.update(k, v.detach(), self.layer_idx,
+                    #                       fmap_key_states=f_k.detach(),
+                    #                       accumulate_in_fp32=True)
+                    past_key_value.update(k, v, self.layer_idx,
+                                          fmap_key_states=f_k,
+                                          accumulate_in_fp32=True)
+            # Concatenate heads and apply output projection
+            _y_true = y_true.transpose(1, 2).contiguous()
+            y_true = self.o_proj(_y_true.view(b, l, self.hidden_size))
+
+            if self.train_attention:
+                attn_weights = _y_true  # flash_attn outputs are shape (b, l, h, d)
+        return y_true, attn_weights, past_key_value
+
+
+class LinearAttentionSlidingWindowCache(LinearAttentionState):
+    """
+    Class for `past_key_values`
+    -> Alternative to KV cache; here we only maintain a "KV state" and "K state"
+    -> Modified from transformers.cache_utils.DynamicCache (v4.36)
+    """
+    def __init__(self, window_size: int = 64) -> None:
+        super().__init__()
+        self._seen_tokens = 0  # should be `self.seen_tokens` in Transformers v4.36
+        self._seen_tokens_by_layer: List[int] = []
+        self.kv_states: List[torch.Tensor] = []
+        self.k_states:  List[torch.Tensor] = []
+
+        # Account for sliding windows
+        self.decode_kv_states: List[torch.Tensor] = []
+        self.decode_k_states: List[torch.Tensor] = []
+        self.k_cache: List[torch.Tensor] = []
+        self.v_cache: List[torch.Tensor] = []
+        self.window_size = window_size
+
+    def update(self, key_states: torch.Tensor, value_states: torch.Tensor, 
+               layer_idx: Optional[int] = None, cache_kwargs: Optional[any] = None,
+               accumulate_in_fp32: bool = False, 
+               fmap_key_states: torch.Tensor = None,  # should not be None
+               grad_enabled: bool = False,
+               **kwargs: any,
+              ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Update KV, K states; and KV cache during training
+        - For decoding, use `self.decode_kv_states` to keep track of KV states 
+          up to sliding window terms
+        - For (chunked) training, use `self.kv_states` to keep track of KV states
+          up to end of sequence
+        - Likewise for `self.decode_k_states` and `self.k_states`
+        """
+        with torch.set_grad_enabled(grad_enabled):
+            if layer_idx == 0:
+                self._seen_tokens += key_states.shape[-2]
+
+            dtype = key_states.dtype
+            if accumulate_in_fp32:
+                # key_states = key_states.float()
+                fmap_key_states = fmap_key_states.float()
+                value_states = value_states.float()
+
+            # Decoding KV state (KV terms up to last window_size)
+            decode_kv_state = torch.einsum(
+                'bhlf,bhld->bhfd', fmap_key_states[:, :, :-self.window_size], value_states[:, :, :-self.window_size]
+            )
+            # KV state
+            kv_state = decode_kv_state + torch.einsum(
+                'bhlf,bhld->bhfd', fmap_key_states[:, :, -self.window_size:], value_states[:, :, -self.window_size:]
+            )
+            # shape is b, h, 1, f; note the 1
+            decode_k_state = fmap_key_states[:, :, :-self.window_size].sum(dim=-2, keepdim=True)
+            k_state = (decode_k_state + fmap_key_states[:, :, -self.window_size:].sum(dim=-2, keepdim=True))
+
+            # Update the cache
+            if len(self.k_states) <= layer_idx:  # Initializing kv and k states
+                self.kv_states.append(kv_state.to(dtype))
+                self.k_states.append(k_state.to(dtype))
+
+                self.decode_kv_states.append(decode_kv_state.to(dtype))
+                self.decode_k_states.append(decode_k_state.to(dtype))
+
+                self.k_cache.append(key_states[:, :, -self.window_size:, :])
+                self.v_cache.append(value_states[:, :, -self.window_size:, :].to(dtype))
+                # self._seen_tokens_by_layer[layer_idx].append(key_states.shape[-2])
+            else:
+                # Update kv and k states recurrently
+                kv_state = (self.kv_states[layer_idx].to(kv_state.dtype) + kv_state).to(dtype)
+                k_state  = (self.k_states[layer_idx].to(kv_state.dtype) + k_state).to(dtype)
+                self.kv_states[layer_idx] = kv_state
+                self.k_states[layer_idx]  = k_state
+
+                decode_kv_state = (self.decode_kv_states[layer_idx].to(kv_state.dtype) 
+                                   + decode_kv_state).to(dtype)
+                decode_k_state  = (self.decode_k_states[layer_idx].to(kv_state.dtype) 
+                                   + decode_k_state).to(dtype)
+                self.decode_kv_states[layer_idx] = decode_kv_state
+                self.decode_k_states[layer_idx]  = decode_k_state
+
+                self.k_cache[layer_idx] = key_states[:, :, -self.window_size:, :]
+                self.v_cache[layer_idx] = value_states[:, :, -self.window_size:, :]
+            self._seen_tokens_by_layer[layer_idx] += key_states.shape[-2]
+
+        return self.kv_states[layer_idx], self.k_states[layer_idx]
+
+    def update_for_decoding(self, keys: torch.Tensor, values: torch.Tensor, 
+                            layer_idx: int, feature_map_k: Callable, dtype: torch.dtype):
+        """
+        Update the decoding KV and K states, and KV cache, during decodeing
+        """
+        with torch.no_grad():
+            k_cache = self.k_cache[layer_idx]
+            v_cache = self.v_cache[layer_idx]
+
+            if k_cache.shape[-2] < self.window_size:  # build window-size cache
+                self.k_cache[layer_idx] = torch.cat([k_cache, keys], dim=-2)
+                self.v_cache[layer_idx] = torch.cat([v_cache, values], dim=-2)
+            else:
+                # MZ 6/3: handle short inputs; zero-out padding when initial k.shape[2] < self.window_size
+                # if k_cache[:, :, :1, :].sum() == 0:   # heuristic for zeroing out padding in cache
+                #     f_k_state = torch.zeros(k_cache[:, :, :1, :].shape, dtype=dtype, device=k_cache.device)
+                # else:
+                #     f_k_state = feature_map_k(k_cache[:, :, :1, :])
+                # -> MZ (later): above only relevant if we zero-pad in our hybrid attention computation
+                k_state = feature_map_k(k_cache[:, :, :1, :])
+                v_state = v_cache[:, :, :1, :]
+                kv_state = torch.einsum('bhlf,bhld->bhfd', k_state.float(), v_state.float()).to(dtype) # b, h, f, d
+                self.decode_kv_states[layer_idx] += kv_state
+                self.decode_k_states[layer_idx] += k_state
+                
+                self.k_cache[layer_idx] = torch.cat([k_cache[:, :, 1:, :], keys], dim=-2)
+                self.v_cache[layer_idx] = torch.cat([v_cache[:, :, 1:, :], values], dim=-2)
+            
+            if layer_idx == 0:
+                self._seen_tokens += keys.shape[-2]
+            self._seen_tokens_by_layer[layer_idx] += keys.shape[-2]
+            return (self.k_cache[layer_idx], self.v_cache[layer_idx], 
+                    self.decode_kv_states[layer_idx], self.decode_k_states[layer_idx])
+
+
+# -----------------
+# Flash Attention 2
+# -----------------
+
+def flash_attention_2(self,
+                      hidden_states: torch.Tensor,
+                      attention_mask: Optional[torch.LongTensor] = None,
+                      position_ids: Optional[torch.LongTensor] = None,
+                      past_key_value: Optional[Cache] = None,
+                      output_attentions: bool = False,
+                      use_cache: bool = False,
+                      cache_position: Optional[torch.LongTensor] = None,
+                     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    """
+    Wrapper for LlamaFlashAttention2
+    Copied and modified from HF Transformers v4.36 and v4.43 implementations
+    - (4.43) https://github.com/huggingface/transformers/blob/868d36d29ec132deeaaf8571b25b6a1b911d0145/src/transformers/models/llama/modeling_llama.py#L402
+    - (4.36) https://github.com/huggingface/transformers/blob/a7cab3c283312b8d4de5df3bbe719971e24f4281/src/transformers/models/llama/modeling_llama.py#L456
+    """
+    output_attentions = False
+
+    bsz, q_len, _ = hidden_states.size()
+
+    query_states = self.q_proj(hidden_states)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)
+
+    # Flash attention requires the input to have the shape
+    # batch_size x seq_length x head_dim x hidden_dim
+    # therefore we just need to keep the original shape
+    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+    try: # As in Transformers v4.36
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
+        cos, sin = self.rotary_emb(key_states, seq_len=kv_seq_len)    
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+    except:  # As in Transformers v4.39
+        cos, sin = self.rotary_emb(key_states, position_ids)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+    if past_key_value is not None:
+        # sin and cos are specific to RoPE models; cache_position needed for the static cache
+        cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+
+    # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+    # to be able to avoid many of these transpose/reshape/view.
+    query_states = query_states.transpose(1, 2)
+    key_states = key_states.transpose(1, 2)
+    value_states = value_states.transpose(1, 2)
+
+    dropout_rate = self.attention_dropout if self.training else 0.0
+
+    # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+    # therefore the input hidden states gets silently casted in float32. Hence, we need
+    # cast them back in the correct dtype just to be sure everything works as expected.
+    # This might slowdown training & inference so it is recommended to not cast the LayerNorms
+    # in fp32. (LlamaRMSNorm handles it correctly)
+
+    input_dtype = query_states.dtype
+    if input_dtype == torch.float32:
+        if torch.is_autocast_enabled():
+            target_dtype = torch.get_autocast_gpu_dtype()
+        # Handle the case where the model is quantized
+        elif hasattr(self.config, "_pre_quantization_dtype"):
+            target_dtype = self.config._pre_quantization_dtype
+        else:
+            target_dtype = self.q_proj.weight.dtype
+
+        logger.warning_once(
+            f"The input hidden states seems to be silently casted in float32, this might be related to"
+            f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+            f" {target_dtype}."
+        )
+
+        query_states = query_states.to(target_dtype)
+        key_states = key_states.to(target_dtype)
+        value_states = value_states.to(target_dtype)
+
+    if getattr(self, '_flash_attention_forward', False):
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate,
+            is_causal=True,
+        )
+    else:
+        attn_output = _flash_attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            q_len,
+            dropout=0, # dropout_rate,
+            sliding_window=getattr(self, "sliding_window", None),
+            use_top_left_mask=self._flash_attn_uses_top_left_mask,
+            is_causal=True,
+        )
+    return attn_output, past_key_value

src/model/linear_attention/linear_window_attention_sw_long.py

@@ -0,0 +1,23 @@
+"""
+LoLCATs attention combining sliding window and linear attentions
+- Using standard sliding window arrangement
+- Training over long sequences with fixed memory with recurrent view
+- During attention transfer, use Flash Attention to compute softmax attention outputs
+
+For each layer: 
+- We first compute (softmax) attention over sliding windows
+- We then compute standard linear attention to "fill in" the earlier parts
+- We combine to model the entire sequence
+"""
+from .linear_window_attention_tk_long import LolcatsTKWindowLongAttention
+from .linear_window_attention_sw import hybrid_attention_quadratic
+
+
+class LolcatsSlidingWindowLongAttention(LolcatsTKWindowLongAttention):
+    """
+    Lolcats attention combining sliding window and linear attention
+    """
+    def __init__(self, remove_base_attn=True, **kwargs):
+        # keep self.base_attn for Flash Attention inference
+        super().__init__(remove_base_attn=True, **kwargs)
+        self.quadratic_attention = hybrid_attention_quadratic