@seanhogg/builderforce-memory-engine 2026.6.18

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (113) hide show
  1. package/LICENSE +21 -0
  2. package/README.md +393 -0
  3. package/dist/index.d.ts +32 -0
  4. package/dist/index.d.ts.map +1 -0
  5. package/dist/index.js +40 -0
  6. package/dist/index.js.map +1 -0
  7. package/dist/kernels/activations.d.ts +5 -0
  8. package/dist/kernels/activations.d.ts.map +1 -0
  9. package/dist/kernels/activations.js +171 -0
  10. package/dist/kernels/activations.js.map +1 -0
  11. package/dist/kernels/attention.d.ts +19 -0
  12. package/dist/kernels/attention.d.ts.map +1 -0
  13. package/dist/kernels/attention.js +263 -0
  14. package/dist/kernels/attention.js.map +1 -0
  15. package/dist/kernels/complex_ssd.d.ts +33 -0
  16. package/dist/kernels/complex_ssd.d.ts.map +1 -0
  17. package/dist/kernels/complex_ssd.js +305 -0
  18. package/dist/kernels/complex_ssd.js.map +1 -0
  19. package/dist/kernels/conv1d.d.ts +3 -0
  20. package/dist/kernels/conv1d.d.ts.map +1 -0
  21. package/dist/kernels/conv1d.js +158 -0
  22. package/dist/kernels/conv1d.js.map +1 -0
  23. package/dist/kernels/linear_projection.d.ts +3 -0
  24. package/dist/kernels/linear_projection.d.ts.map +1 -0
  25. package/dist/kernels/linear_projection.js +219 -0
  26. package/dist/kernels/linear_projection.js.map +1 -0
  27. package/dist/kernels/selective_scan.d.ts +3 -0
  28. package/dist/kernels/selective_scan.d.ts.map +1 -0
  29. package/dist/kernels/selective_scan.js +348 -0
  30. package/dist/kernels/selective_scan.js.map +1 -0
  31. package/dist/kernels/ssd.d.ts +29 -0
  32. package/dist/kernels/ssd.d.ts.map +1 -0
  33. package/dist/kernels/ssd.js +276 -0
  34. package/dist/kernels/ssd.js.map +1 -0
  35. package/dist/kernels/weight_update.d.ts +3 -0
  36. package/dist/kernels/weight_update.d.ts.map +1 -0
  37. package/dist/kernels/weight_update.js +119 -0
  38. package/dist/kernels/weight_update.js.map +1 -0
  39. package/dist/model/attention_block.d.ts +48 -0
  40. package/dist/model/attention_block.d.ts.map +1 -0
  41. package/dist/model/attention_block.js +262 -0
  42. package/dist/model/attention_block.js.map +1 -0
  43. package/dist/model/mamba1_block.d.ts +70 -0
  44. package/dist/model/mamba1_block.d.ts.map +1 -0
  45. package/dist/model/mamba1_block.js +333 -0
  46. package/dist/model/mamba1_block.js.map +1 -0
  47. package/dist/model/mamba2_block.d.ts +44 -0
  48. package/dist/model/mamba2_block.d.ts.map +1 -0
  49. package/dist/model/mamba2_block.js +252 -0
  50. package/dist/model/mamba2_block.js.map +1 -0
  51. package/dist/model/mamba3_block.d.ts +51 -0
  52. package/dist/model/mamba3_block.d.ts.map +1 -0
  53. package/dist/model/mamba3_block.js +270 -0
  54. package/dist/model/mamba3_block.js.map +1 -0
  55. package/dist/model/mamba_block.d.ts +64 -0
  56. package/dist/model/mamba_block.d.ts.map +1 -0
  57. package/dist/model/mamba_block.js +303 -0
  58. package/dist/model/mamba_block.js.map +1 -0
  59. package/dist/model/mamba_model.d.ts +140 -0
  60. package/dist/model/mamba_model.d.ts.map +1 -0
  61. package/dist/model/mamba_model.js +527 -0
  62. package/dist/model/mamba_model.js.map +1 -0
  63. package/dist/model/sequence_layer.d.ts +25 -0
  64. package/dist/model/sequence_layer.d.ts.map +1 -0
  65. package/dist/model/sequence_layer.js +8 -0
  66. package/dist/model/sequence_layer.js.map +1 -0
  67. package/dist/tokenizer/bpe.d.ts +29 -0
  68. package/dist/tokenizer/bpe.d.ts.map +1 -0
  69. package/dist/tokenizer/bpe.js +164 -0
  70. package/dist/tokenizer/bpe.js.map +1 -0
  71. package/dist/training/autograd.d.ts +27 -0
  72. package/dist/training/autograd.d.ts.map +1 -0
  73. package/dist/training/autograd.js +120 -0
  74. package/dist/training/autograd.js.map +1 -0
  75. package/dist/training/trainer.d.ts +36 -0
  76. package/dist/training/trainer.d.ts.map +1 -0
  77. package/dist/training/trainer.js +183 -0
  78. package/dist/training/trainer.js.map +1 -0
  79. package/dist/utils/gpu_utils.d.ts +21 -0
  80. package/dist/utils/gpu_utils.d.ts.map +1 -0
  81. package/dist/utils/gpu_utils.js +111 -0
  82. package/dist/utils/gpu_utils.js.map +1 -0
  83. package/dist/utils/quantization.d.ts +26 -0
  84. package/dist/utils/quantization.d.ts.map +1 -0
  85. package/dist/utils/quantization.js +116 -0
  86. package/dist/utils/quantization.js.map +1 -0
  87. package/dist/utils/rng.d.ts +36 -0
  88. package/dist/utils/rng.d.ts.map +1 -0
  89. package/dist/utils/rng.js +61 -0
  90. package/dist/utils/rng.js.map +1 -0
  91. package/package.json +99 -0
  92. package/src/index.ts +114 -0
  93. package/src/kernels/activations.ts +174 -0
  94. package/src/kernels/attention.ts +268 -0
  95. package/src/kernels/complex_ssd.ts +307 -0
  96. package/src/kernels/conv1d.ts +159 -0
  97. package/src/kernels/linear_projection.ts +220 -0
  98. package/src/kernels/selective_scan.ts +350 -0
  99. package/src/kernels/ssd.ts +278 -0
  100. package/src/kernels/weight_update.ts +120 -0
  101. package/src/model/attention_block.ts +344 -0
  102. package/src/model/mamba1_block.ts +437 -0
  103. package/src/model/mamba2_block.ts +319 -0
  104. package/src/model/mamba3_block.ts +335 -0
  105. package/src/model/mamba_block.ts +401 -0
  106. package/src/model/mamba_model.ts +678 -0
  107. package/src/model/sequence_layer.ts +29 -0
  108. package/src/tokenizer/bpe.ts +186 -0
  109. package/src/training/autograd.ts +135 -0
  110. package/src/training/trainer.ts +309 -0
  111. package/src/utils/gpu_utils.ts +147 -0
  112. package/src/utils/quantization.ts +154 -0
  113. package/src/utils/rng.ts +65 -0
@@ -0,0 +1 @@
1
+ {"version":3,"file":"mamba2_block.js","sourceRoot":"","sources":["../../src/model/mamba2_block.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;GAWG;AAEH,OAAO,EACH,qBAAqB,EACrB,eAAe,EACf,mBAAmB,EACnB,wBAAwB,EACxB,mBAAmB,EACnB,cAAc,EACd,IAAI,GACP,MAAM,uBAAuB,CAAC;AAE/B,OAAO,EAAE,gBAAgB,EAAE,MAAS,mBAAmB,CAAC;AACxD,OAAO,EAAE,aAAa,EAAE,MAAM,iBAAiB,CAAC;AAChD,OAAO,EAAE,mBAAmB,EAAE,MAAM,sBAAsB,CAAC;AAC3D,OAAO,EAAE,mBAAmB,EAAE,MAAM,iCAAiC,CAAC;AACtE,OAAO,EAAE,gBAAgB,EAAE,MAAS,2BAA2B,CAAC;AAkBhE,MAAM,UAAU,GAAG,UAAU,CAAA;;;;;;;;;;CAU5B,CAAC;AAEF,MAAM,OAAO,WAAW;IACX,SAAS,GAAG,QAAiB,CAAC;IAEvC,MAAM,CAAa;IACnB,MAAM,CAA+B;IACrC,MAAM,CAAU;IAChB,KAAK,CAAW;IAEhB,UAAU,CAA6B;IACvC,SAAS,CAAuC;IAExC,SAAS,GAAG,KAAK,CAAC;IAE1B,YAAY,MAAiB,EAAE,MAAyB;QACpD,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC;QACrB,IAAI,CAAC,MAAM,GAAG;YACV,GAAG,EAAE,MAAM,EAAE,EAAE,EAAE,KAAK,EAAE,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,OAAO,EAAE,CAAC,EAAE,QAAQ,EAAE,GAAG,EAAE;YACjE,GAAG,MAAM;SACmB,CAAC;QAEjC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QAC/C,IAAI,CAAC,MAAM,GAAG,MAAM,GAAG,MAAM,CAAC;QAC9B,IAAI,CAAC,KAAK,GAAI,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC;QAEnC,IAAI,IAAI,CAAC,MAAM,GAAG,MAAM,KAAK,CAAC,EAAE,CAAC;YAC7B,MAAM,IAAI,KAAK,CACX,wBAAwB,IAAI,CAAC,MAAM,kCAAkC,MAAM,IAAI,CAClF,CAAC;QACN,CAAC;QAED,IAAI,CAAC,UAAU,GAAG,EAAE,CAAC;QACrB,IAAI,CAAC,SAAS,GAAI,EAAE,CAAC;QAErB,IAAI,CAAC,YAAY,EAAE,CAAC;QACpB,IAAI,CAAC,eAAe,EAAE,CAAC;IAC3B,CAAC;IAEO,YAAY;QAChB,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,OAAO,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QAC/D,MAAM,CAAC,GAAI,IAAI,CAAC,MAAM,CAAC;QACvB,MAAM,CAAC,GAAI,MAAM,CAAC;QAClB,MAAM,CAAC,GAAI,KAAK,CAAC;QACjB,MAAM,CAAC,GAAI,MAAM,CAAC;QAClB,MAAM,CAAC,GAAI,OAAO,CAAC;QAEnB,MAAM,KAAK,GAAG,CAAC,CAAS,EAAE,GAAG,GAAG,IAAI,EAAgB,EAAE,CAAC,aAAa,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC;QAE7E,MAAM,KAAK,GAAG,CAAC,CAAS,EAAE,EAAE,CAAC,IAAI,YAAY,CAAC,CAAC,CAAC,CAAC;QACjD,MAAM,IAAI,GAAI,CAAC,CAAS,EAAE,EAAE,CAAC,IAAI,YAAY,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;QAE3D,4EAA4E;QAC5E,MAAM,UAAU,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC;QACrC,MAAM,EAAE,GAAG,CAAC,GAAiB,EAAE,EAAE,CAAC,mBAAmB,CAAC,IAAI,CAAC,MAAM,EAAE,GAAG,EAAE,IAAI,CAAC,CAAC;QAE9E,IAAI,CAAC,UAAU,GAAG;YACd,OAAO,EAAO,EAAE,CAAC,KAAK,CAAC,UAAU,GAAG,MAAM,CAAC,CAAC;YAC5C,KAAK,EAAS,EAAE,CAAC,KAAK,CAAC,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;YAClD,KAAK,EAAS,EAAE,CAAC,KAAK,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;YACtC,KAAK,EAAS,EAAE,CAAC,IAAI,YAAY,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;YACzD,OAAO,EAAO,EAAE,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;YAC1B,KAAK,EAAS,EAAE,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;YACzB,QAAQ,EAAM,EAAE,CAAC,KAAK,CAAC,MAAM,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;YACzC,UAAU,EAAI,EAAE,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,EAAW,gBAAgB;YACpD,aAAa,EAAE,EAAE,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC,EAAK,oBAAoB;SAC3D,CAAC;IACN,CAAC;IAEO,eAAe;QACnB,MAAM,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC;QACtB,IAAI,CAAC,SAAS,GAAG;YACb,MAAM,EAAK,qBAAqB,CAAC,CAAC,EAAE,mBAAmB,EAAG,gBAAgB,CAAC;YAC3E,MAAM,EAAK,qBAAqB,CAAC,CAAC,EAAE,mBAAmB,EAAG,gBAAgB,CAAC;YAC3E,OAAO,EAAI,qBAAqB,CAAC,CAAC,EAAE,gBAAgB,EAAM,iBAAiB,CAAC;YAC5E,OAAO,EAAI,qBAAqB,CAAC,CAAC,EAAE,gBAAgB,EAAM,mBAAmB,CAAC;YAC9E,KAAK,EAAM,qBAAqB,CAAC,CAAC,EAAE,UAAU,EAAY,MAAM,CAAC;SACpE,CAAC;IACN,CAAC;IAED,OAAO,CAAC,IAAe,EAAE,KAAa,EAAE,MAAc;QAClD,MAAM,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC;QACtB,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,OAAO,EAAE,QAAQ,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QACzE,MAAM,CAAC,GAAI,IAAI,CAAC,MAAM,CAAC;QACvB,MAAM,CAAC,GAAI,MAAM,CAAC;QAClB,MAAM,CAAC,GAAI,KAAK,CAAC;QACjB,MAAM,CAAC,GAAI,MAAM,CAAC;QAClB,MAAM,CAAC,GAAI,OAAO,CAAC;QACnB,MAAM,EAAE,GAAG,IAAI,CAAC,KAAK,CAAC;QACtB,MAAM,CAAC,GAAI,KAAK,CAAC;QACjB,MAAM,CAAC,GAAI,MAAM,CAAC;QAClB,MAAM,CAAC,GAAI,CAAC,GAAG,CAAC,CAAC;QACjB,MAAM,KAAK,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAE,wCAAwC;QACtE,MAAM,SAAS,GAAG,IAAI,CAAC,IAAI,CAAC,CAAC,GAAG,QAAQ,CAAC,CAAC;QAE1C,uBAAuB;QACvB,MAAM,OAAO,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,MAAM,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAClE,MAAM,OAAO,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACzD,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,EAAE,CAAC,CAAC;YACnC,IAAI,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC;YAC/C,IAAI,YAAY,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC;YAC3C,MAAM,IAAI,GAAG,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC5C,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,SAAS,CAAE,EACpD,CAAC,IAAI,EAAE,IAAI,EAAE,IAAI,CAAC,UAAU,CAAC,eAAe,CAAE,EAAE,OAAO,EAAE,OAAO,CAAC,CAAC,CAAC;YACvE,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,SAAS,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAC3E,CAAC;QACD,OAAO,CAAC,OAAO,EAAE,CAAC;QAElB,iEAAiE;QACjE,MAAM,UAAU,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC;QACrC,MAAM,SAAS,GAAI,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,UAAU,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACzE,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,UAAU,CAAC,CAAC,CAAC,MAAM,CAAC;YAC/D,MAAM,IAAI,GAAK,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC9C,kDAAkD;YAClD,MAAM,QAAQ,GAAG,mBAAmB,CAAC,CAAC,EAAE,IAAI,YAAY,CAAC,UAAU,CAAC,EAAE,IAAI,CAAC,CAAC;YAC5E,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EACnD,CAAC,IAAI,EAAE,OAAO,EAAE,IAAI,CAAC,UAAU,CAAC,SAAS,CAAE,EAAE,QAAQ,EAAE,SAAS,CAAC,CAAC,CAAC;YACvE,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,IAAI,CAAC,UAAU,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;YACzF,QAAQ,CAAC,OAAO,EAAE,CAAC;QACvB,CAAC;QACD,OAAO,CAAC,OAAO,EAAE,CAAC;QAElB,+BAA+B;QAC/B,MAAM,QAAQ,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,KAAK,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAClE,MAAM,KAAK,GAAM,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAC9D,CAAC;YACG,MAAM,GAAG,GAAG,CAAC,CAAC,oBAAoB,EAAE,CAAC;YACrC,GAAG,CAAC,kBAAkB,CAAC,SAAS,EAAE,CAAC,EAAe,QAAQ,EAAE,CAAC,EAAE,CAAC,GAAG,KAAK,GAAG,CAAC,CAAC,CAAC;YAC9E,GAAG,CAAC,kBAAkB,CAAC,SAAS,EAAE,CAAC,GAAG,KAAK,GAAG,CAAC,EAAG,KAAK,EAAK,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;YAC1E,CAAC,CAAC,KAAK,CAAC,MAAM,CAAC,CAAC,GAAG,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC;QACnC,CAAC;QACD,SAAS,CAAC,OAAO,EAAE,CAAC;QAEpB,oEAAoE;QACpE,MAAM,OAAO,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,KAAK,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACjE,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,CAAC,CAAC,EAAE,KAAK,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC;YAC3D,MAAM,IAAI,GAAK,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC9C,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EACnD,CAAC,IAAI,EAAE,QAAQ,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAE,OAAO,CAAC,CAAC,CAAC;YACrF,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,IAAI,CAAC,KAAK,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACxF,CAAC;QACD,QAAQ,CAAC,OAAO,EAAE,CAAC;QAEnB,uDAAuD;QACvD,MAAM,OAAO,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAC7D,MAAM,QAAQ,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAClE,MAAM,QAAQ,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAClE,CAAC;YACG,MAAM,GAAG,GAAG,CAAC,CAAC,oBAAoB,EAAE,CAAC;YACrC,GAAG,CAAC,kBAAkB,CAAC,OAAO,EAAE,CAAC,EAAa,OAAO,EAAG,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;YACtE,GAAG,CAAC,kBAAkB,CAAC,OAAO,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAK,QAAQ,EAAE,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;YAC1E,GAAG,CAAC,kBAAkB,CAAC,OAAO,EAAE,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;YACjF,CAAC,CAAC,KAAK,CAAC,MAAM,CAAC,CAAC,GAAG,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC;QACnC,CAAC;QACD,OAAO,CAAC,OAAO,EAAE,CAAC;QAElB,cAAc;QACd,6CAA6C;QAC7C,MAAM,UAAU,GAAG,wBAAwB,CACvC,CAAC,EAAE,CAAC,SAAS,GAAG,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACnD,MAAM,MAAM,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAE5D,CAAC;YACG,MAAM,SAAS,GAAG,IAAI,WAAW,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,QAAQ,EAAE,SAAS,EAAE,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC;YACtF,MAAM,IAAI,GAAG,mBAAmB,CAAC,CAAC,EAAE,SAAS,CAAC,CAAC;YAC/C,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,SAAS,CAAE,EACpD,CAAC,IAAI,EAAE,OAAO,EAAE,QAAQ,EAAE,QAAQ,EAAE,KAAK;gBACxC,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAE,IAAI,CAAC,UAAU,CAAC,SAAS,CAAE;gBACtD,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAE,MAAM,EAAE,UAAU,CAAC,CAAC,CAAC;YACrD,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,SAAS,CAAE,EAAE,EAAE,EAAE,CAAC,SAAS,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACzE,CAAC;QACD,OAAO,CAAC,OAAO,EAAE,CAAC;QAClB,QAAQ,CAAC,OAAO,EAAE,CAAC;QACnB,QAAQ,CAAC,OAAO,EAAE,CAAC;QACnB,KAAK,CAAC,OAAO,EAAE,CAAC;QAEhB,kCAAkC;QAClC,MAAM,YAAY,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAClE,MAAM,YAAY,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAC9D,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,EAAE,CAAC,CAAC;YACnC,IAAI,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;YAC1C,IAAI,YAAY,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC;YAC3C,MAAM,IAAI,GAAG,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC5C,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,SAAS,CAAE,EACpD,CAAC,IAAI,EAAE,MAAM,EAAE,IAAI,CAAC,UAAU,CAAC,YAAY,CAAE,EAAE,YAAY,EAAE,YAAY,CAAC,CAAC,CAAC;YAChF,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,SAAS,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAC3E,CAAC;QACD,MAAM,CAAC,OAAO,EAAE,CAAC;QACjB,YAAY,CAAC,OAAO,EAAE,CAAC;QAEvB,uBAAuB;QACvB,MAAM,UAAU,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,MAAM,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACrE,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,MAAM,CAAC;YACtD,MAAM,IAAI,GAAK,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC9C,MAAM,QAAQ,GAAG,mBAAmB,CAAC,CAAC,EAAE,IAAI,YAAY,CAAC,MAAM,CAAC,EAAE,IAAI,CAAC,CAAC;YACxE,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EACnD,CAAC,IAAI,EAAE,YAAY,EAAE,IAAI,CAAC,UAAU,CAAC,UAAU,CAAE,EAAE,QAAQ,EAAE,UAAU,CAAC,CAAC,CAAC;YAC9E,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,IAAI,CAAC,MAAM,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;YACrF,QAAQ,CAAC,OAAO,EAAE,CAAC;QACvB,CAAC;QACD,YAAY,CAAC,OAAO,EAAE,CAAC;QAEvB,kBAAkB;QAClB,MAAM,MAAM,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,MAAM,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACjE,CAAC;YACG,MAAM,IAAI,GAAG,mBAAmB,CAAC,CAAC,EAAE,IAAI,WAAW,CAAC,CAAC,CAAC,GAAG,MAAM,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC;YAC1E,MAAM,EAAE,GAAK,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,OAAO,CAAE,EACpD,CAAC,UAAU,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,CAAC,CAAC,CAAC;YACtC,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,OAAO,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,GAAG,MAAM,EAAE,GAAG,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACnF,CAAC;QACD,UAAU,CAAC,OAAO,EAAE,CAAC;QAErB,MAAM,KAAK,GAAgB,EAAE,UAAU,EAAE,CAAC;QAC1C,OAAO,EAAE,MAAM,EAAE,KAAK,EAAE,CAAC;IAC7B,CAAC;IAED,UAAU;QACN,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,OAAO,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QAC/D,MAAM,CAAC,GAAO,IAAI,CAAC,MAAM,CAAC;QAC1B,MAAM,CAAC,GAAO,MAAM,CAAC;QACrB,MAAM,CAAC,GAAO,KAAK,CAAC;QACpB,MAAM,CAAC,GAAO,MAAM,CAAC;QACrB,MAAM,CAAC,GAAO,OAAO,CAAC;QACtB,MAAM,KAAK,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC;QAE5B,OAAO;YACH,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,SAAS,CAAE,EAAO,KAAK,EAAE,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,GAAG,MAAM,EAAE,IAAI,EAAE,SAAS,EAAO;YACpG,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAU,KAAK,EAAE,KAAK,GAAG,CAAC,EAAqB,IAAI,EAAE,OAAO,EAAS;YACrG,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAU,KAAK,EAAE,KAAK,EAAyB,IAAI,EAAE,OAAO,EAAS;YACrG,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAU,KAAK,EAAE,CAAC,EAA6B,IAAI,EAAE,OAAO,EAAS;YACrG,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,SAAS,CAAE,EAAQ,KAAK,EAAE,CAAC,EAA6B,IAAI,EAAE,SAAS,EAAO;YACrG,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAU,KAAK,EAAE,CAAC,EAA6B,IAAI,EAAE,OAAO,EAAS;YACrG,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,UAAU,CAAE,EAAO,KAAK,EAAE,MAAM,GAAG,CAAC,EAAoB,IAAI,EAAE,UAAU,EAAM;YACrG,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,YAAY,CAAE,EAAK,KAAK,EAAE,CAAC,EAA6B,IAAI,EAAE,YAAY,EAAI;YACrG,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,eAAe,CAAE,EAAE,KAAK,EAAE,MAAM,EAAwB,IAAI,EAAE,eAAe,EAAC;SACxG,CAAC;IACN,CAAC;IAED,kBAAkB;QACd,IAAI,IAAI,CAAC,SAAS,EAAE,CAAC;YACjB,uEAAuE;YACvE,OAAO;gBACH,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,SAAS,CAAE;oBAChC,KAAK,EAAE,CAAC,IAAI,CAAC,MAAM,CAAC,OAAO,GAAG,IAAI,CAAC,MAAM,CAAC,MAAM,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC,MAAM;oBAC1E,IAAI,EAAE,YAAY,EAAE;aACzB,CAAC;QACN,CAAC;QACD,OAAO,IAAI,CAAC,UAAU,EAAE,CAAC;IAC7B,CAAC;IAED,WAAW,CAAC,OAAgB;QACxB,IAAI,CAAC,SAAS,GAAG,OAAO,CAAC;IAC7B,CAAC;IAED,OAAO;QACH,KAAK,MAAM,GAAG,IAAI,MAAM,CAAC,MAAM,CAAC,IAAI,CAAC,UAAU,CAAC;YAAE,GAAG,CAAC,OAAO,EAAE,CAAC;QAChE,IAAI,CAAC,UAAU,GAAG,EAAE,CAAC;IACzB,CAAC;CACJ"}
@@ -0,0 +1,51 @@
1
+ /**
2
+ * mamba3_block.ts – Mamba-3 Mixer Block (Complex-valued MIMO SSM, inference-first).
3
+ *
4
+ * Three improvements over Mamba-2:
5
+ * 1. Complex-valued states — h ∈ ℂ^(N/2), stored as interleaved f32 pairs
6
+ * 2. MIMO recurrence — G×G block recurrence per head (default G=1 = SISO)
7
+ * 3. ET discretisation — B_bar = (A_bar − 1)·A⁻¹·B (exact, not approx)
8
+ *
9
+ * Weight shapes vs Mamba-2 (same 9 tensors, different A_log shape):
10
+ * wInProj : (D + 2*G*N_c*2 + H, dModel) where N_c = dState (complex count)
11
+ * wConv : (D + 2*G*N_c*2, K)
12
+ * bConv : (D + 2*G*N_c*2,)
13
+ * A_log : (H, 2) ← [log|A|, arg(A)] per head
14
+ * dt_bias : (H,)
15
+ * D_vec : (H,)
16
+ * wOutProj : (dModel, D)
17
+ * normWeight : (D,)
18
+ * preNormWeight: (dModel,)
19
+ *
20
+ * Implements SequenceLayer.
21
+ */
22
+ import type { Mamba2BlockConfig } from './mamba2_block.js';
23
+ import type { SequenceLayer, LayerForwardResult, LayerParam } from './sequence_layer.js';
24
+ export interface Mamba3BlockConfig extends Mamba2BlockConfig {
25
+ /** MIMO group size G. Default 1 = SISO (same as Mamba-2). */
26
+ mimoGroup?: number;
27
+ }
28
+ export interface Mamba3Cache {
29
+ stateCarry: GPUBuffer;
30
+ }
31
+ export declare class Mamba3Block implements SequenceLayer {
32
+ readonly layerType: "mamba3";
33
+ device: GPUDevice;
34
+ config: Required<Mamba3BlockConfig>;
35
+ dInner: number;
36
+ dHead: number;
37
+ /** Complex state count per head (N_c = dState in config). */
38
+ nComplex: number;
39
+ gpuWeights: Record<string, GPUBuffer>;
40
+ pipelines: Record<string, GPUComputePipeline>;
41
+ private _wslaMode;
42
+ constructor(device: GPUDevice, config: Mamba3BlockConfig);
43
+ private _initWeights;
44
+ private _buildPipelines;
45
+ forward(xBuf: GPUBuffer, batch: number, seqLen: number): LayerForwardResult;
46
+ parameters(): LayerParam[];
47
+ getTrainableParams(): LayerParam[];
48
+ setWSLAMode(enabled: boolean): void;
49
+ destroy(): void;
50
+ }
51
+ //# sourceMappingURL=mamba3_block.d.ts.map
@@ -0,0 +1 @@
1
+ {"version":3,"file":"mamba3_block.d.ts","sourceRoot":"","sources":["../../src/model/mamba3_block.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;GAoBG;AAkBH,OAAO,KAAK,EAAE,iBAAiB,EAAE,MAAuB,mBAAmB,CAAC;AAC5E,OAAO,KAAK,EAAE,aAAa,EAAE,kBAAkB,EAAE,UAAU,EAAE,MAAM,qBAAqB,CAAC;AAEzF,MAAM,WAAW,iBAAkB,SAAQ,iBAAiB;IACxD,6DAA6D;IAC7D,SAAS,CAAC,EAAE,MAAM,CAAC;CAEtB;AAED,MAAM,WAAW,WAAW;IACxB,UAAU,EAAE,SAAS,CAAC;CACzB;AAcD,qBAAa,WAAY,YAAW,aAAa;IAC7C,QAAQ,CAAC,SAAS,EAAG,QAAQ,CAAU;IAEvC,MAAM,EAAM,SAAS,CAAC;IACtB,MAAM,EAAM,QAAQ,CAAC,iBAAiB,CAAC,CAAC;IACxC,MAAM,EAAM,MAAM,CAAC;IACnB,KAAK,EAAO,MAAM,CAAC;IACnB,6DAA6D;IAC7D,QAAQ,EAAI,MAAM,CAAC;IAEnB,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,SAAS,CAAC,CAAC;IACtC,SAAS,EAAG,MAAM,CAAC,MAAM,EAAE,kBAAkB,CAAC,CAAC;IAE/C,OAAO,CAAC,SAAS,CAAS;gBAEd,MAAM,EAAE,SAAS,EAAE,MAAM,EAAE,iBAAiB;IAyBxD,OAAO,CAAC,YAAY;IAuCpB,OAAO,CAAC,eAAe;IAWvB,OAAO,CAAC,IAAI,EAAE,SAAS,EAAE,KAAK,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,kBAAkB;IA4I3E,UAAU,IAAI,UAAU,EAAE;IAsB1B,kBAAkB,IAAI,UAAU,EAAE;IAWlC,WAAW,CAAC,OAAO,EAAE,OAAO,GAAG,IAAI;IAInC,OAAO,IAAI,IAAI;CAIlB"}
@@ -0,0 +1,270 @@
1
+ /**
2
+ * mamba3_block.ts – Mamba-3 Mixer Block (Complex-valued MIMO SSM, inference-first).
3
+ *
4
+ * Three improvements over Mamba-2:
5
+ * 1. Complex-valued states — h ∈ ℂ^(N/2), stored as interleaved f32 pairs
6
+ * 2. MIMO recurrence — G×G block recurrence per head (default G=1 = SISO)
7
+ * 3. ET discretisation — B_bar = (A_bar − 1)·A⁻¹·B (exact, not approx)
8
+ *
9
+ * Weight shapes vs Mamba-2 (same 9 tensors, different A_log shape):
10
+ * wInProj : (D + 2*G*N_c*2 + H, dModel) where N_c = dState (complex count)
11
+ * wConv : (D + 2*G*N_c*2, K)
12
+ * bConv : (D + 2*G*N_c*2,)
13
+ * A_log : (H, 2) ← [log|A|, arg(A)] per head
14
+ * dt_bias : (H,)
15
+ * D_vec : (H,)
16
+ * wOutProj : (dModel, D)
17
+ * normWeight : (D,)
18
+ * preNormWeight: (dModel,)
19
+ *
20
+ * Implements SequenceLayer.
21
+ */
22
+ import { createComputePipeline, createBindGroup, createStorageBuffer, createEmptyStorageBuffer, createUniformBuffer, dispatchKernel, cdiv, } from '../utils/gpu_utils.js';
23
+ import { COMPLEX_SSD_FORWARD_WGSL } from '../kernels/complex_ssd.js';
24
+ import { gaussianArray } from '../utils/rng.js';
25
+ import { CONV1D_FORWARD_WGSL } from '../kernels/conv1d.js';
26
+ import { LINEAR_FORWARD_WGSL } from '../kernels/linear_projection.js';
27
+ import { ACTIVATIONS_WGSL } from '../kernels/activations.js';
28
+ const ADD_SHADER = /* wgsl */ `
29
+ @group(0) @binding(0) var<storage, read> a : array<f32>;
30
+ @group(0) @binding(1) var<storage, read> b : array<f32>;
31
+ @group(0) @binding(2) var<storage, read_write> c : array<f32>;
32
+ @group(0) @binding(3) var<uniform> n : u32;
33
+ @compute @workgroup_size(256)
34
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
35
+ let i = gid.x;
36
+ if (i < n) { c[i] = a[i] + b[i]; }
37
+ }
38
+ `;
39
+ export class Mamba3Block {
40
+ layerType = 'mamba3';
41
+ device;
42
+ config;
43
+ dInner;
44
+ dHead;
45
+ /** Complex state count per head (N_c = dState in config). */
46
+ nComplex;
47
+ gpuWeights;
48
+ pipelines;
49
+ _wslaMode = false;
50
+ constructor(device, config) {
51
+ this.device = device;
52
+ this.config = {
53
+ ...{ dState: 16, dConv: 4, expand: 2, nGroups: 1, chunkLen: 256, mimoGroup: 1 },
54
+ ...config,
55
+ };
56
+ const { dModel, expand, nHeads } = this.config;
57
+ this.dInner = expand * dModel;
58
+ this.dHead = this.dInner / nHeads;
59
+ this.nComplex = this.config.dState; // N_c
60
+ if (this.dInner % nHeads !== 0) {
61
+ throw new Error(`Mamba3Block: dInner (${this.dInner}) must be divisible by nHeads (${nHeads}).`);
62
+ }
63
+ this.gpuWeights = {};
64
+ this.pipelines = {};
65
+ this._initWeights();
66
+ this._buildPipelines();
67
+ }
68
+ _initWeights() {
69
+ const { dModel, dConv, nHeads, nGroups } = this.config;
70
+ const D = this.dInner;
71
+ const Nc = this.nComplex;
72
+ const K = dConv;
73
+ const H = nHeads;
74
+ const G = nGroups;
75
+ // Each complex state = 2 f32 values
76
+ const convD = D + 2 * G * Nc * 2; // x-channels + complex B/C
77
+ const randn = (n, std = 0.02) => gaussianArray(n, std);
78
+ const zeros = (n) => new Float32Array(n);
79
+ const ones = (n) => new Float32Array(n).fill(1.0);
80
+ // A_log: (H, 2) = [log|A|, arg(A)] per head
81
+ // Initialise to unit magnitude (|A|=1, phase=0) → purely oscillatory
82
+ const A_log = new Float32Array(H * 2);
83
+ for (let h = 0; h < H; h++) {
84
+ A_log[h * 2 + 0] = 0.0; // log|A| = 0 → |A| = 1
85
+ A_log[h * 2 + 1] = (2 * Math.PI * h) / H; // evenly spaced phases
86
+ }
87
+ const mk = (arr) => createStorageBuffer(this.device, arr, true);
88
+ const inProjRows = D + 2 * G * Nc * 2 + H;
89
+ this.gpuWeights = {
90
+ wInProj: mk(randn(inProjRows * dModel)),
91
+ wConv: mk(randn(convD * K, 0.01)),
92
+ bConv: mk(zeros(convD)),
93
+ A_log: mk(A_log),
94
+ dt_bias: mk(zeros(H)),
95
+ D_vec: mk(ones(H)),
96
+ wOutProj: mk(randn(dModel * D, 0.02)),
97
+ normWeight: mk(ones(D)),
98
+ preNormWeight: mk(ones(dModel)),
99
+ };
100
+ }
101
+ _buildPipelines() {
102
+ const d = this.device;
103
+ this.pipelines = {
104
+ linear: createComputePipeline(d, LINEAR_FORWARD_WGSL, 'linear_forward'),
105
+ conv1d: createComputePipeline(d, CONV1D_FORWARD_WGSL, 'conv1d_forward'),
106
+ rmsnorm: createComputePipeline(d, ACTIVATIONS_WGSL, 'rmsnorm_forward'),
107
+ cssd_fwd: createComputePipeline(d, COMPLEX_SSD_FORWARD_WGSL, 'complex_ssd_forward'),
108
+ elAdd: createComputePipeline(d, ADD_SHADER, 'main'),
109
+ };
110
+ }
111
+ forward(xBuf, batch, seqLen) {
112
+ const d = this.device;
113
+ const { dModel, dConv, nHeads, nGroups, chunkLen } = this.config;
114
+ const D = this.dInner;
115
+ const Nc = this.nComplex;
116
+ const K = dConv;
117
+ const H = nHeads;
118
+ const G = nGroups;
119
+ const dh = this.dHead;
120
+ const B = batch;
121
+ const L = seqLen;
122
+ const M = B * L;
123
+ const convD = D + 2 * G * Nc * 2;
124
+ const numChunks = Math.ceil(L / chunkLen);
125
+ // 1. Pre-block RMSNorm
126
+ const normOut = createEmptyStorageBuffer(d, M * dModel * 4, true);
127
+ const normInv = createEmptyStorageBuffer(d, M * 4, true);
128
+ {
129
+ const params = new ArrayBuffer(16);
130
+ new Uint32Array(params, 0, 2).set([M, dModel]);
131
+ new Float32Array(params, 8, 1).set([1e-6]);
132
+ const pBuf = createUniformBuffer(d, params);
133
+ const bg = createBindGroup(d, this.pipelines['rmsnorm'], [pBuf, xBuf, this.gpuWeights['preNormWeight'], normOut, normInv]);
134
+ dispatchKernel(d, this.pipelines['rmsnorm'], bg, [cdiv(M, 64), 1, 1]);
135
+ }
136
+ normInv.destroy();
137
+ // 2. Fused in_proj
138
+ const inProjRows = D + 2 * G * Nc * 2 + H;
139
+ const inProjOut = createEmptyStorageBuffer(d, M * inProjRows * 4, true);
140
+ {
141
+ const params = new Uint32Array([M, dModel, inProjRows]).buffer;
142
+ const pBuf = createUniformBuffer(d, params);
143
+ const zeroBias = createStorageBuffer(d, new Float32Array(inProjRows), true);
144
+ const bg = createBindGroup(d, this.pipelines['linear'], [pBuf, normOut, this.gpuWeights['wInProj'], zeroBias, inProjOut]);
145
+ dispatchKernel(d, this.pipelines['linear'], bg, [cdiv(M, 16), cdiv(inProjRows, 16), 1]);
146
+ zeroBias.destroy();
147
+ }
148
+ normOut.destroy();
149
+ // Split: xConv [convD], dt [H]
150
+ const xConvBuf = createEmptyStorageBuffer(d, M * convD * 4, true);
151
+ const dtBuf = createEmptyStorageBuffer(d, M * H * 4, true);
152
+ {
153
+ const enc = d.createCommandEncoder();
154
+ enc.copyBufferToBuffer(inProjOut, 0, xConvBuf, 0, M * convD * 4);
155
+ enc.copyBufferToBuffer(inProjOut, M * convD * 4, dtBuf, 0, M * H * 4);
156
+ d.queue.submit([enc.finish()]);
157
+ }
158
+ inProjOut.destroy();
159
+ // 3. Causal conv1d (fused convD channels)
160
+ const convOut = createEmptyStorageBuffer(d, M * convD * 4, true);
161
+ {
162
+ const params = new Uint32Array([L, convD, K, B, 1]).buffer;
163
+ const pBuf = createUniformBuffer(d, params);
164
+ const bg = createBindGroup(d, this.pipelines['conv1d'], [pBuf, xConvBuf, this.gpuWeights['wConv'], this.gpuWeights['bConv'], convOut]);
165
+ dispatchKernel(d, this.pipelines['conv1d'], bg, [cdiv(L, 16), cdiv(convD, 16), B]);
166
+ }
167
+ xConvBuf.destroy();
168
+ // Split: xSsd [D], B_proj_complex [G*Nc*2], C_proj_complex [G*Nc*2]
169
+ const xSsdBuf = createEmptyStorageBuffer(d, M * D * 4, true);
170
+ const bProjBuf = createEmptyStorageBuffer(d, M * G * Nc * 2 * 4, true);
171
+ const cProjBuf = createEmptyStorageBuffer(d, M * G * Nc * 2 * 4, true);
172
+ {
173
+ const enc = d.createCommandEncoder();
174
+ enc.copyBufferToBuffer(convOut, 0, xSsdBuf, 0, M * D * 4);
175
+ enc.copyBufferToBuffer(convOut, M * D * 4, bProjBuf, 0, M * G * Nc * 2 * 4);
176
+ enc.copyBufferToBuffer(convOut, M * (D + G * Nc * 2) * 4, cProjBuf, 0, M * G * Nc * 2 * 4);
177
+ d.queue.submit([enc.finish()]);
178
+ }
179
+ convOut.destroy();
180
+ // 4. Complex SSD scan
181
+ // state_carry: [numChunks+1, B, H, Nc*2, dHead]
182
+ const stateCarry = createEmptyStorageBuffer(d, (numChunks + 1) * B * H * Nc * 2 * dh * 4, true);
183
+ const cssdOut = createEmptyStorageBuffer(d, M * D * 4, true);
184
+ {
185
+ const params = new Uint32Array([L, D, H, dh, G, Nc, chunkLen, numChunks, B]).buffer;
186
+ const pBuf = createUniformBuffer(d, params);
187
+ const bg = createBindGroup(d, this.pipelines['cssd_fwd'], [pBuf, xSsdBuf, bProjBuf, cProjBuf, dtBuf,
188
+ this.gpuWeights['A_log'], this.gpuWeights['dt_bias'],
189
+ this.gpuWeights['D_vec'], cssdOut, stateCarry]);
190
+ dispatchKernel(d, this.pipelines['cssd_fwd'], bg, [numChunks, H, B]);
191
+ }
192
+ xSsdBuf.destroy();
193
+ bProjBuf.destroy();
194
+ cProjBuf.destroy();
195
+ dtBuf.destroy();
196
+ // 5. Inner RMSNorm
197
+ const innerNormOut = createEmptyStorageBuffer(d, M * D * 4, true);
198
+ const innerNormInv = createEmptyStorageBuffer(d, M * 4, true);
199
+ {
200
+ const params = new ArrayBuffer(16);
201
+ new Uint32Array(params, 0, 2).set([M, D]);
202
+ new Float32Array(params, 8, 1).set([1e-6]);
203
+ const pBuf = createUniformBuffer(d, params);
204
+ const bg = createBindGroup(d, this.pipelines['rmsnorm'], [pBuf, cssdOut, this.gpuWeights['normWeight'], innerNormOut, innerNormInv]);
205
+ dispatchKernel(d, this.pipelines['rmsnorm'], bg, [cdiv(M, 64), 1, 1]);
206
+ }
207
+ cssdOut.destroy();
208
+ innerNormInv.destroy();
209
+ // 6. Output projection
210
+ const outProjOut = createEmptyStorageBuffer(d, M * dModel * 4, true);
211
+ {
212
+ const params = new Uint32Array([M, D, dModel]).buffer;
213
+ const pBuf = createUniformBuffer(d, params);
214
+ const zeroBias = createStorageBuffer(d, new Float32Array(dModel), true);
215
+ const bg = createBindGroup(d, this.pipelines['linear'], [pBuf, innerNormOut, this.gpuWeights['wOutProj'], zeroBias, outProjOut]);
216
+ dispatchKernel(d, this.pipelines['linear'], bg, [cdiv(M, 16), cdiv(dModel, 16), 1]);
217
+ zeroBias.destroy();
218
+ }
219
+ innerNormOut.destroy();
220
+ // 7. Residual add
221
+ const output = createEmptyStorageBuffer(d, M * dModel * 4, true);
222
+ {
223
+ const nBuf = createUniformBuffer(d, new Uint32Array([M * dModel]).buffer);
224
+ const bg = createBindGroup(d, this.pipelines['elAdd'], [outProjOut, xBuf, output, nBuf]);
225
+ dispatchKernel(d, this.pipelines['elAdd'], bg, [cdiv(M * dModel, 256), 1, 1]);
226
+ }
227
+ outProjOut.destroy();
228
+ const cache = { stateCarry };
229
+ return { output, cache };
230
+ }
231
+ parameters() {
232
+ const { dModel, dConv, nHeads, nGroups } = this.config;
233
+ const D = this.dInner;
234
+ const Nc = this.nComplex;
235
+ const K = dConv;
236
+ const H = nHeads;
237
+ const G = nGroups;
238
+ const convD = D + 2 * G * Nc * 2;
239
+ return [
240
+ { buf: this.gpuWeights['wInProj'], numel: (D + 2 * G * Nc * 2 + H) * dModel, name: 'wInProj' },
241
+ { buf: this.gpuWeights['wConv'], numel: convD * K, name: 'wConv' },
242
+ { buf: this.gpuWeights['bConv'], numel: convD, name: 'bConv' },
243
+ { buf: this.gpuWeights['A_log'], numel: H * 2, name: 'A_log' },
244
+ { buf: this.gpuWeights['dt_bias'], numel: H, name: 'dt_bias' },
245
+ { buf: this.gpuWeights['D_vec'], numel: H, name: 'D_vec' },
246
+ { buf: this.gpuWeights['wOutProj'], numel: dModel * D, name: 'wOutProj' },
247
+ { buf: this.gpuWeights['normWeight'], numel: D, name: 'normWeight' },
248
+ { buf: this.gpuWeights['preNormWeight'], numel: dModel, name: 'preNormWeight' },
249
+ ];
250
+ }
251
+ getTrainableParams() {
252
+ if (this._wslaMode) {
253
+ return [
254
+ { buf: this.gpuWeights['wInProj'],
255
+ numel: (this.config.nGroups * this.nComplex * 2 * 2) * this.config.dModel,
256
+ name: 'wInProj_BC' },
257
+ ];
258
+ }
259
+ return this.parameters();
260
+ }
261
+ setWSLAMode(enabled) {
262
+ this._wslaMode = enabled;
263
+ }
264
+ destroy() {
265
+ for (const buf of Object.values(this.gpuWeights))
266
+ buf.destroy();
267
+ this.gpuWeights = {};
268
+ }
269
+ }
270
+ //# sourceMappingURL=mamba3_block.js.map
@@ -0,0 +1 @@
1
+ {"version":3,"file":"mamba3_block.js","sourceRoot":"","sources":["../../src/model/mamba3_block.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;GAoBG;AAEH,OAAO,EACH,qBAAqB,EACrB,eAAe,EACf,mBAAmB,EACnB,wBAAwB,EACxB,mBAAmB,EACnB,cAAc,EACd,IAAI,GACP,MAAM,uBAAuB,CAAC;AAE/B,OAAO,EAAE,wBAAwB,EAAE,MAAM,2BAA2B,CAAC;AACrE,OAAO,EAAE,aAAa,EAAE,MAAM,iBAAiB,CAAC;AAChD,OAAO,EAAE,mBAAmB,EAAE,MAAW,sBAAsB,CAAC;AAChE,OAAO,EAAE,mBAAmB,EAAE,MAAW,iCAAiC,CAAC;AAC3E,OAAO,EAAE,gBAAgB,EAAE,MAAc,2BAA2B,CAAC;AAerE,MAAM,UAAU,GAAG,UAAU,CAAA;;;;;;;;;;CAU5B,CAAC;AAEF,MAAM,OAAO,WAAW;IACX,SAAS,GAAG,QAAiB,CAAC;IAEvC,MAAM,CAAgB;IACtB,MAAM,CAAkC;IACxC,MAAM,CAAa;IACnB,KAAK,CAAc;IACnB,6DAA6D;IAC7D,QAAQ,CAAW;IAEnB,UAAU,CAA4B;IACtC,SAAS,CAAsC;IAEvC,SAAS,GAAG,KAAK,CAAC;IAE1B,YAAY,MAAiB,EAAE,MAAyB;QACpD,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC;QACrB,IAAI,CAAC,MAAM,GAAG;YACV,GAAG,EAAE,MAAM,EAAE,EAAE,EAAE,KAAK,EAAE,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,OAAO,EAAE,CAAC,EAAE,QAAQ,EAAE,GAAG,EAAE,SAAS,EAAE,CAAC,EAAE;YAC/E,GAAG,MAAM;SACmB,CAAC;QAEjC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QAC/C,IAAI,CAAC,MAAM,GAAK,MAAM,GAAG,MAAM,CAAC;QAChC,IAAI,CAAC,KAAK,GAAM,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC;QACrC,IAAI,CAAC,QAAQ,GAAG,IAAI,CAAC,MAAM,CAAC,MAAM,CAAC,CAAC,MAAM;QAE1C,IAAI,IAAI,CAAC,MAAM,GAAG,MAAM,KAAK,CAAC,EAAE,CAAC;YAC7B,MAAM,IAAI,KAAK,CACX,wBAAwB,IAAI,CAAC,MAAM,kCAAkC,MAAM,IAAI,CAClF,CAAC;QACN,CAAC;QAED,IAAI,CAAC,UAAU,GAAG,EAAE,CAAC;QACrB,IAAI,CAAC,SAAS,GAAI,EAAE,CAAC;QAErB,IAAI,CAAC,YAAY,EAAE,CAAC;QACpB,IAAI,CAAC,eAAe,EAAE,CAAC;IAC3B,CAAC;IAEO,YAAY;QAChB,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,OAAO,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QACvD,MAAM,CAAC,GAAI,IAAI,CAAC,MAAM,CAAC;QACvB,MAAM,EAAE,GAAG,IAAI,CAAC,QAAQ,CAAC;QACzB,MAAM,CAAC,GAAI,KAAK,CAAC;QACjB,MAAM,CAAC,GAAI,MAAM,CAAC;QAClB,MAAM,CAAC,GAAI,OAAO,CAAC;QACnB,oCAAoC;QACpC,MAAM,KAAK,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,CAAC,CAAE,2BAA2B;QAE9D,MAAM,KAAK,GAAG,CAAC,CAAS,EAAE,GAAG,GAAG,IAAI,EAAgB,EAAE,CAAC,aAAa,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC;QAE7E,MAAM,KAAK,GAAG,CAAC,CAAS,EAAE,EAAE,CAAC,IAAI,YAAY,CAAC,CAAC,CAAC,CAAC;QACjD,MAAM,IAAI,GAAI,CAAC,CAAS,EAAE,EAAE,CAAC,IAAI,YAAY,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC;QAE3D,4CAA4C;QAC5C,qEAAqE;QACrE,MAAM,KAAK,GAAG,IAAI,YAAY,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC;QACtC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;YACzB,KAAK,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,GAAG,GAAG,CAAC,CAAyB,uBAAuB;YACvE,KAAK,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,EAAE,GAAG,CAAC,CAAC,GAAG,CAAC,CAAC,CAAM,uBAAuB;QAC1E,CAAC;QAED,MAAM,EAAE,GAAG,CAAC,GAAiB,EAAE,EAAE,CAAC,mBAAmB,CAAC,IAAI,CAAC,MAAM,EAAE,GAAG,EAAE,IAAI,CAAC,CAAC;QAE9E,MAAM,UAAU,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,GAAG,CAAC,CAAC;QAC1C,IAAI,CAAC,UAAU,GAAG;YACd,OAAO,EAAO,EAAE,CAAC,KAAK,CAAC,UAAU,GAAG,MAAM,CAAC,CAAC;YAC5C,KAAK,EAAS,EAAE,CAAC,KAAK,CAAC,KAAK,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;YACxC,KAAK,EAAS,EAAE,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC;YAC9B,KAAK,EAAS,EAAE,CAAC,KAAK,CAAC;YACvB,OAAO,EAAO,EAAE,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;YAC1B,KAAK,EAAS,EAAE,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;YACzB,QAAQ,EAAM,EAAE,CAAC,KAAK,CAAC,MAAM,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;YACzC,UAAU,EAAI,EAAE,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;YACzB,aAAa,EAAE,EAAE,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;SAClC,CAAC;IACN,CAAC;IAEO,eAAe;QACnB,MAAM,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC;QACtB,IAAI,CAAC,SAAS,GAAG;YACb,MAAM,EAAO,qBAAqB,CAAC,CAAC,EAAE,mBAAmB,EAAQ,gBAAgB,CAAC;YAClF,MAAM,EAAO,qBAAqB,CAAC,CAAC,EAAE,mBAAmB,EAAQ,gBAAgB,CAAC;YAClF,OAAO,EAAM,qBAAqB,CAAC,CAAC,EAAE,gBAAgB,EAAW,iBAAiB,CAAC;YACnF,QAAQ,EAAK,qBAAqB,CAAC,CAAC,EAAE,wBAAwB,EAAG,qBAAqB,CAAC;YACvF,KAAK,EAAQ,qBAAqB,CAAC,CAAC,EAAE,UAAU,EAAiB,MAAM,CAAC;SAC3E,CAAC;IACN,CAAC;IAED,OAAO,CAAC,IAAe,EAAE,KAAa,EAAE,MAAc;QAClD,MAAM,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC;QACtB,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,OAAO,EAAE,QAAQ,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QACjE,MAAM,CAAC,GAAI,IAAI,CAAC,MAAM,CAAC;QACvB,MAAM,EAAE,GAAG,IAAI,CAAC,QAAQ,CAAC;QACzB,MAAM,CAAC,GAAI,KAAK,CAAC;QACjB,MAAM,CAAC,GAAI,MAAM,CAAC;QAClB,MAAM,CAAC,GAAI,OAAO,CAAC;QACnB,MAAM,EAAE,GAAG,IAAI,CAAC,KAAK,CAAC;QACtB,MAAM,CAAC,GAAI,KAAK,CAAC;QACjB,MAAM,CAAC,GAAI,MAAM,CAAC;QAClB,MAAM,CAAC,GAAI,CAAC,GAAG,CAAC,CAAC;QACjB,MAAM,KAAK,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,CAAC;QACjC,MAAM,SAAS,GAAG,IAAI,CAAC,IAAI,CAAC,CAAC,GAAG,QAAQ,CAAC,CAAC;QAE1C,uBAAuB;QACvB,MAAM,OAAO,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,MAAM,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAClE,MAAM,OAAO,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACzD,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,EAAE,CAAC,CAAC;YACnC,IAAI,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC;YAC/C,IAAI,YAAY,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC;YAC3C,MAAM,IAAI,GAAG,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC5C,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,SAAS,CAAE,EACpD,CAAC,IAAI,EAAE,IAAI,EAAE,IAAI,CAAC,UAAU,CAAC,eAAe,CAAE,EAAE,OAAO,EAAE,OAAO,CAAC,CAAC,CAAC;YACvE,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,SAAS,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAC3E,CAAC;QACD,OAAO,CAAC,OAAO,EAAE,CAAC;QAElB,mBAAmB;QACnB,MAAM,UAAU,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,GAAG,CAAC,CAAC;QAC1C,MAAM,SAAS,GAAI,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,UAAU,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACzE,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,CAAC,CAAC,EAAE,MAAM,EAAE,UAAU,CAAC,CAAC,CAAC,MAAM,CAAC;YAC/D,MAAM,IAAI,GAAK,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC9C,MAAM,QAAQ,GAAG,mBAAmB,CAAC,CAAC,EAAE,IAAI,YAAY,CAAC,UAAU,CAAC,EAAE,IAAI,CAAC,CAAC;YAC5E,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EACnD,CAAC,IAAI,EAAE,OAAO,EAAE,IAAI,CAAC,UAAU,CAAC,SAAS,CAAE,EAAE,QAAQ,EAAE,SAAS,CAAC,CAAC,CAAC;YACvE,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,IAAI,CAAC,UAAU,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;YACzF,QAAQ,CAAC,OAAO,EAAE,CAAC;QACvB,CAAC;QACD,OAAO,CAAC,OAAO,EAAE,CAAC;QAElB,+BAA+B;QAC/B,MAAM,QAAQ,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,KAAK,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAClE,MAAM,KAAK,GAAM,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAC9D,CAAC;YACG,MAAM,GAAG,GAAG,CAAC,CAAC,oBAAoB,EAAE,CAAC;YACrC,GAAG,CAAC,kBAAkB,CAAC,SAAS,EAAE,CAAC,EAAc,QAAQ,EAAE,CAAC,EAAE,CAAC,GAAG,KAAK,GAAG,CAAC,CAAC,CAAC;YAC7E,GAAG,CAAC,kBAAkB,CAAC,SAAS,EAAE,CAAC,GAAG,KAAK,GAAG,CAAC,EAAE,KAAK,EAAK,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;YACzE,CAAC,CAAC,KAAK,CAAC,MAAM,CAAC,CAAC,GAAG,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC;QACnC,CAAC;QACD,SAAS,CAAC,OAAO,EAAE,CAAC;QAEpB,0CAA0C;QAC1C,MAAM,OAAO,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,KAAK,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACjE,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,CAAC,CAAC,EAAE,KAAK,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC;YAC3D,MAAM,IAAI,GAAK,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC9C,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EACnD,CAAC,IAAI,EAAE,QAAQ,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAE,OAAO,CAAC,CAAC,CAAC;YACrF,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,IAAI,CAAC,KAAK,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACxF,CAAC;QACD,QAAQ,CAAC,OAAO,EAAE,CAAC;QAEnB,oEAAoE;QACpE,MAAM,OAAO,GAAI,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAC9D,MAAM,QAAQ,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACvE,MAAM,QAAQ,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACvE,CAAC;YACG,MAAM,GAAG,GAAG,CAAC,CAAC,oBAAoB,EAAE,CAAC;YACrC,GAAG,CAAC,kBAAkB,CAAC,OAAO,EAAE,CAAC,EAAwB,OAAO,EAAG,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;YACjF,GAAG,CAAC,kBAAkB,CAAC,OAAO,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAgB,QAAQ,EAAE,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;YAC1F,GAAG,CAAC,kBAAkB,CAAC,OAAO,EAAE,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,CAAC,GAAG,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;YAC3F,CAAC,CAAC,KAAK,CAAC,MAAM,CAAC,CAAC,GAAG,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC;QACnC,CAAC;QACD,OAAO,CAAC,OAAO,EAAE,CAAC;QAElB,sBAAsB;QACtB,gDAAgD;QAChD,MAAM,UAAU,GAAG,wBAAwB,CACvC,CAAC,EAAE,CAAC,SAAS,GAAG,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACxD,MAAM,OAAO,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAE7D,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,EAAE,EAAE,CAAC,EAAE,EAAE,EAAE,QAAQ,EAAE,SAAS,EAAE,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC;YACpF,MAAM,IAAI,GAAK,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC9C,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,UAAU,CAAE,EACrD,CAAC,IAAI,EAAE,OAAO,EAAE,QAAQ,EAAE,QAAQ,EAAE,KAAK;gBACxC,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAE,IAAI,CAAC,UAAU,CAAC,SAAS,CAAE;gBACtD,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAE,OAAO,EAAE,UAAU,CAAC,CAAC,CAAC;YACtD,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,UAAU,CAAE,EAAE,EAAE,EAAE,CAAC,SAAS,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAC1E,CAAC;QACD,OAAO,CAAC,OAAO,EAAE,CAAC;QAClB,QAAQ,CAAC,OAAO,EAAE,CAAC;QACnB,QAAQ,CAAC,OAAO,EAAE,CAAC;QACnB,KAAK,CAAC,OAAO,EAAE,CAAC;QAEhB,mBAAmB;QACnB,MAAM,YAAY,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAClE,MAAM,YAAY,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QAC9D,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,EAAE,CAAC,CAAC;YACnC,IAAI,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;YAC1C,IAAI,YAAY,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC;YAC3C,MAAM,IAAI,GAAG,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC5C,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,SAAS,CAAE,EACpD,CAAC,IAAI,EAAE,OAAO,EAAE,IAAI,CAAC,UAAU,CAAC,YAAY,CAAE,EAAE,YAAY,EAAE,YAAY,CAAC,CAAC,CAAC;YACjF,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,SAAS,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAC3E,CAAC;QACD,OAAO,CAAC,OAAO,EAAE,CAAC;QAClB,YAAY,CAAC,OAAO,EAAE,CAAC;QAEvB,uBAAuB;QACvB,MAAM,UAAU,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,MAAM,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACrE,CAAC;YACG,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,MAAM,CAAC;YACtD,MAAM,IAAI,GAAK,mBAAmB,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YAC9C,MAAM,QAAQ,GAAG,mBAAmB,CAAC,CAAC,EAAE,IAAI,YAAY,CAAC,MAAM,CAAC,EAAE,IAAI,CAAC,CAAC;YACxE,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EACnD,CAAC,IAAI,EAAE,YAAY,EAAE,IAAI,CAAC,UAAU,CAAC,UAAU,CAAE,EAAE,QAAQ,EAAE,UAAU,CAAC,CAAC,CAAC;YAC9E,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,EAAE,EAAE,CAAC,EAAE,IAAI,CAAC,MAAM,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;YACrF,QAAQ,CAAC,OAAO,EAAE,CAAC;QACvB,CAAC;QACD,YAAY,CAAC,OAAO,EAAE,CAAC;QAEvB,kBAAkB;QAClB,MAAM,MAAM,GAAG,wBAAwB,CAAC,CAAC,EAAE,CAAC,GAAG,MAAM,GAAG,CAAC,EAAE,IAAI,CAAC,CAAC;QACjE,CAAC;YACG,MAAM,IAAI,GAAG,mBAAmB,CAAC,CAAC,EAAE,IAAI,WAAW,CAAC,CAAC,CAAC,GAAG,MAAM,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC;YAC1E,MAAM,EAAE,GAAK,eAAe,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,OAAO,CAAE,EACpD,CAAC,UAAU,EAAE,IAAI,EAAE,MAAM,EAAE,IAAI,CAAC,CAAC,CAAC;YACtC,cAAc,CAAC,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,OAAO,CAAE,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,GAAG,MAAM,EAAE,GAAG,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QACnF,CAAC;QACD,UAAU,CAAC,OAAO,EAAE,CAAC;QAErB,MAAM,KAAK,GAAgB,EAAE,UAAU,EAAE,CAAC;QAC1C,OAAO,EAAE,MAAM,EAAE,KAAK,EAAE,CAAC;IAC7B,CAAC;IAED,UAAU;QACN,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,OAAO,EAAE,GAAG,IAAI,CAAC,MAAM,CAAC;QACvD,MAAM,CAAC,GAAO,IAAI,CAAC,MAAM,CAAC;QAC1B,MAAM,EAAE,GAAM,IAAI,CAAC,QAAQ,CAAC;QAC5B,MAAM,CAAC,GAAO,KAAK,CAAC;QACpB,MAAM,CAAC,GAAO,MAAM,CAAC;QACrB,MAAM,CAAC,GAAO,OAAO,CAAC;QACtB,MAAM,KAAK,GAAG,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,CAAC;QAEjC,OAAO;YACH,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,SAAS,CAAE,EAAO,KAAK,EAAE,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,GAAG,CAAC,CAAC,GAAG,MAAM,EAAE,IAAI,EAAE,SAAS,EAAO;YACzG,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAU,KAAK,EAAE,KAAK,GAAG,CAAC,EAA2B,IAAI,EAAE,OAAO,EAAS;YAC3G,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAU,KAAK,EAAE,KAAK,EAA+B,IAAI,EAAE,OAAO,EAAS;YAC3G,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAU,KAAK,EAAE,CAAC,GAAG,CAAC,EAA+B,IAAI,EAAE,OAAO,EAAS;YAC3G,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,SAAS,CAAE,EAAQ,KAAK,EAAE,CAAC,EAAmC,IAAI,EAAE,SAAS,EAAO;YAC3G,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,OAAO,CAAE,EAAU,KAAK,EAAE,CAAC,EAAmC,IAAI,EAAE,OAAO,EAAS;YAC3G,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,UAAU,CAAE,EAAO,KAAK,EAAE,MAAM,GAAG,CAAC,EAA0B,IAAI,EAAE,UAAU,EAAM;YAC3G,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,YAAY,CAAE,EAAK,KAAK,EAAE,CAAC,EAAmC,IAAI,EAAE,YAAY,EAAI;YAC3G,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,eAAe,CAAE,EAAE,KAAK,EAAE,MAAM,EAA8B,IAAI,EAAE,eAAe,EAAC;SAC9G,CAAC;IACN,CAAC;IAED,kBAAkB;QACd,IAAI,IAAI,CAAC,SAAS,EAAE,CAAC;YACjB,OAAO;gBACH,EAAE,GAAG,EAAE,IAAI,CAAC,UAAU,CAAC,SAAS,CAAE;oBAChC,KAAK,EAAE,CAAC,IAAI,CAAC,MAAM,CAAC,OAAO,GAAG,IAAI,CAAC,QAAQ,GAAG,CAAC,GAAG,CAAC,CAAC,GAAG,IAAI,CAAC,MAAM,CAAC,MAAM;oBACzE,IAAI,EAAE,YAAY,EAAE;aACzB,CAAC;QACN,CAAC;QACD,OAAO,IAAI,CAAC,UAAU,EAAE,CAAC;IAC7B,CAAC;IAED,WAAW,CAAC,OAAgB;QACxB,IAAI,CAAC,SAAS,GAAG,OAAO,CAAC;IAC7B,CAAC;IAED,OAAO;QACH,KAAK,MAAM,GAAG,IAAI,MAAM,CAAC,MAAM,CAAC,IAAI,CAAC,UAAU,CAAC;YAAE,GAAG,CAAC,OAAO,EAAE,CAAC;QAChE,IAAI,CAAC,UAAU,GAAG,EAAE,CAAC;IACzB,CAAC;CACJ"}
@@ -0,0 +1,64 @@
1
+ /**
2
+ * mamba_block.ts – Mamba Mixer Block
3
+ */
4
+ export interface MambaBlockConfig {
5
+ dModel: number;
6
+ dState?: number;
7
+ dConv?: number;
8
+ expand?: number;
9
+ dtRank?: number;
10
+ biasConv?: boolean;
11
+ }
12
+ export interface BlockParam {
13
+ buf: GPUBuffer;
14
+ numel: number;
15
+ name: string;
16
+ }
17
+ export interface BlockCache {
18
+ normInv: GPUBuffer;
19
+ normIn: GPUBuffer;
20
+ normOut: GPUBuffer;
21
+ zBuf: GPUBuffer;
22
+ xConvIn: GPUBuffer;
23
+ convOut: GPUBuffer;
24
+ siluOut: GPUBuffer;
25
+ deltaFull: GPUBuffer;
26
+ B_raw: GPUBuffer;
27
+ C_raw: GPUBuffer;
28
+ hCache: GPUBuffer;
29
+ }
30
+ export interface BlockForwardResult {
31
+ output: GPUBuffer;
32
+ cache: BlockCache;
33
+ }
34
+ export declare class MambaBlock {
35
+ device: GPUDevice;
36
+ config: Required<MambaBlockConfig>;
37
+ dInner: number;
38
+ dtRank: number;
39
+ wInProj: Float32Array;
40
+ bInProj: Float32Array;
41
+ wConv: Float32Array;
42
+ bConv: Float32Array;
43
+ wXProj: Float32Array;
44
+ bXProj: Float32Array;
45
+ wDtProj: Float32Array;
46
+ bDtProj: Float32Array;
47
+ A_log: Float32Array;
48
+ D_vec: Float32Array;
49
+ wOutProj: Float32Array;
50
+ bOutProj: Float32Array;
51
+ normWeight: Float32Array;
52
+ gpuWeights: Record<string, GPUBuffer>;
53
+ pipelines: Record<string, GPUComputePipeline>;
54
+ private _wslaMode;
55
+ constructor(device: GPUDevice, config: MambaBlockConfig);
56
+ private _initWeights;
57
+ private _uploadWeightsToGPU;
58
+ private _buildPipelines;
59
+ forward(xBuf: GPUBuffer, batch: number, seqLen: number): BlockForwardResult;
60
+ parameters(): BlockParam[];
61
+ setWSLAMode(enabled: boolean): void;
62
+ getTrainableParams(): BlockParam[];
63
+ }
64
+ //# sourceMappingURL=mamba_block.d.ts.map
@@ -0,0 +1 @@
1
+ {"version":3,"file":"mamba_block.d.ts","sourceRoot":"","sources":["../../src/model/mamba_block.ts"],"names":[],"mappings":"AAAA;;GAEG;AAkBH,MAAM,WAAW,gBAAgB;IAC/B,MAAM,EAAE,MAAM,CAAC;IACf,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,QAAQ,CAAC,EAAE,OAAO,CAAC;CACpB;AAED,MAAM,WAAW,UAAU;IACzB,GAAG,EAAE,SAAS,CAAC;IACf,KAAK,EAAE,MAAM,CAAC;IACd,IAAI,EAAE,MAAM,CAAC;CACd;AAED,MAAM,WAAW,UAAU;IACzB,OAAO,EAAE,SAAS,CAAC;IACnB,MAAM,EAAE,SAAS,CAAC;IAClB,OAAO,EAAE,SAAS,CAAC;IACnB,IAAI,EAAE,SAAS,CAAC;IAChB,OAAO,EAAE,SAAS,CAAC;IACnB,OAAO,EAAE,SAAS,CAAC;IACnB,OAAO,EAAE,SAAS,CAAC;IACnB,SAAS,EAAE,SAAS,CAAC;IACrB,KAAK,EAAE,SAAS,CAAC;IACjB,KAAK,EAAE,SAAS,CAAC;IACjB,MAAM,EAAE,SAAS,CAAC;CACnB;AAED,MAAM,WAAW,kBAAkB;IACjC,MAAM,EAAE,SAAS,CAAC;IAClB,KAAK,EAAE,UAAU,CAAC;CACnB;AAED,qBAAa,UAAU;IACnB,MAAM,EAAE,SAAS,CAAC;IAClB,MAAM,EAAE,QAAQ,CAAC,gBAAgB,CAAC,CAAC;IACnC,MAAM,EAAE,MAAM,CAAC;IACf,MAAM,EAAE,MAAM,CAAC;IACf,OAAO,EAAE,YAAY,CAAC;IACtB,OAAO,EAAE,YAAY,CAAC;IACtB,KAAK,EAAE,YAAY,CAAC;IACpB,KAAK,EAAE,YAAY,CAAC;IACpB,MAAM,EAAE,YAAY,CAAC;IACrB,MAAM,EAAE,YAAY,CAAC;IACrB,OAAO,EAAE,YAAY,CAAC;IACtB,OAAO,EAAE,YAAY,CAAC;IACtB,KAAK,EAAE,YAAY,CAAC;IACpB,KAAK,EAAE,YAAY,CAAC;IACpB,QAAQ,EAAE,YAAY,CAAC;IACvB,QAAQ,EAAE,YAAY,CAAC;IACvB,UAAU,EAAE,YAAY,CAAC;IACzB,UAAU,EAAE,MAAM,CAAC,MAAM,EAAE,SAAS,CAAC,CAAC;IACtC,SAAS,EAAE,MAAM,CAAC,MAAM,EAAE,kBAAkB,CAAC,CAAC;IAC9C,OAAO,CAAC,SAAS,CAAS;gBAEd,MAAM,EAAE,SAAS,EAAE,MAAM,EAAE,gBAAgB;IAoCvD,OAAO,CAAC,YAAY;IAoCpB,OAAO,CAAC,mBAAmB;IAqB3B,OAAO,CAAC,eAAe;IAavB,OAAO,CAAC,IAAI,EAAE,SAAS,EAAE,KAAK,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,kBAAkB;IAqL3E,UAAU,IAAI,UAAU,EAAE;IAwB1B,WAAW,CAAC,OAAO,EAAE,OAAO,GAAG,IAAI;IAInC,kBAAkB,IAAI,UAAU,EAAE;CASrC"}