perf(optimizer): stream factored AdamW4bit state

areno/accel/csrc/attention.cu

@@ -11,7 +11,6 @@
 namespace {
 
 constexpr int kAttentionThreads = 256;
-constexpr int kAttentionTileN = 16;
 
 int64_t attention_tile_n(int64_t head_dim) {
   if (head_dim <= 256) {

areno/accel/csrc/extension.cpp

@@ -192,8 +192,10 @@ void areno_adamw_4bit_step_cuda(
     torch::Tensor exp_avg_scale,
     torch::Tensor exp_avg_sq_q,
     torch::Tensor exp_avg_sq_scale,
-    int64_t packed_offset,
-    int64_t scale_offset,
+    int64_t moment_packed_offset,
+    int64_t moment_scale_offset,
+    int64_t variance_packed_offset,
+    int64_t variance_scale_offset,
     int64_t quant_block_size,
     double beta1,
     double beta2,
@@ -202,6 +204,33 @@ void areno_adamw_4bit_step_cuda(
     double eps,
     double step_size,
     double bias_correction2_sqrt);
+void areno_adamw_4bit_factored_stats_cuda(
+    torch::Tensor grad,
+    torch::Tensor factor_sums,
+    torch::Tensor invalid,
+    int64_t parameter_shard_start,
+    int64_t rows,
+    int64_t columns);
+void areno_adamw_4bit_factored_step_cuda(
+    torch::Tensor model,
+    torch::Tensor grad,
+    torch::Tensor exp_avg_q,
+    torch::Tensor exp_avg_scale,
+    torch::Tensor factors,
+    torch::Tensor row_mean,
+    torch::Tensor invalid,
+    int64_t moment_packed_offset,
+    int64_t moment_scale_offset,
+    int64_t parameter_shard_start,
+    int64_t quant_block_size,
+    int64_t rows,
+    int64_t columns,
+    double beta1,
+    double effective_lr,
+    double weight_decay,
+    double eps,
+    double step_size,
+    double bias_correction2_sqrt);
 void areno_adamw_8bit_step_cuda(
     torch::Tensor model,
     torch::Tensor grad,
@@ -235,6 +264,14 @@ void areno_adamw_fp32_state_step_cuda(
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("areno_adamw_fp32_master_step", &areno_adamw_fp32_master_step_cuda, "ARENO compact FP32-master AdamW step");
   m.def("areno_adamw_4bit_step", &areno_adamw_4bit_step_cuda, "ARENO packed block-wise AdamW4bit step");
+  m.def(
+      "areno_adamw_4bit_factored_stats",
+      &areno_adamw_4bit_factored_stats_cuda,
+      "ARENO factored AdamW4bit statistics pass");
+  m.def(
+      "areno_adamw_4bit_factored_step",
+      &areno_adamw_4bit_factored_step_cuda,
+      "ARENO factored AdamW4bit update pass");
   m.def("areno_adamw_8bit_step", &areno_adamw_8bit_step_cuda, "ARENO block-wise 8-bit AdamW step");
   m.def("areno_adamw_fp32_state_step", &areno_adamw_fp32_state_step_cuda, "ARENO FP32-state AdamW step");
   m.def("areno_silu_and_mul", &areno_silu_and_mul_cuda, "ARENO SiLU and multiply");